<commit_before>"""Base translators for translating Grow content."""

import copy
import json
import logging
import os
import threading
import progressbar
import texttable
import yaml
from protorpc import message_types
from protorpc import messages
from protorpc import protojson
from grow.common import progressbar_non
from grow.common import utils
from grow.translators import errors as translator_errors

class TranslatorStat(messages.Message):
    lang = messages.StringField(1)
    num_words = messages.IntegerField(2)
    num_words_translated = messages.IntegerField(3)
    source_lang = messages.StringField(4)
    ident = messages.StringField(5)
    url = messages.StringField(6)
    uploaded = message_types.DateTimeField(7)
    service = messages.StringField(8)
    downloaded = message_types.DateTimeField(9)

def translator_stat_representer(dumper, stat):
    content = json.loads(protojson.encode_message(stat))
    content.pop('lang')  # Exclude from serialization.
    return dumper.represent_mapping('tag:yaml.org,2002:map', content)

yaml.SafeDumper.add_representer(TranslatorStat,
                                translator_stat_representer)

class TranslatorServiceError(Exception):

    def __init__(self, message, ident=None, locale=None):
        if locale:
            new_message = 'Error for locale "{}" -> {}'.format(locale, message)
        elif ident:
            new_message = 'Error for resource "{}" -> {}'.format(
                ident, message)
        else:
            new_message = message
        super(TranslatorServiceError, self).__init__(new_message)

class Translator(object):
    TRANSLATOR_STATS_PATH = '/translators.yaml'
    KIND = None
    has_immutable_translation_resources = False
    has_multiple_langs_in_one_resource = False

    def __init__(self, pod, config=None, project_title=None,
                 instructions=None, inject=False):
        self.pod = pod
        self.config = config or {}
        self.project_title = project_title or 'Untitled Grow Website'
        self.instructions = instructions
        self._inject = inject

    def _cleanup_locales(self, locales):
        """Certain locales should be ignored."""
        clean_locales = []
        default_locale = self.pod.podspec.default_locale
        skipped = {
            'symlink': set(),
            'po': set(),
        }
        for locale in locales:
            locale_path = os.path.join('translations', str(locale))

            # Silently ignore the default locale.
            if default_locale and str(locale) == str(default_locale):
                continue

            # Ignore the symlinks.
            if os.path.islink(locale_path):
                skipped['symlink'].add(str(locale))
                continue

            # Ignore the locales without a `.PO` file.
            po_path = os.path.join(locale_path, 'LC_MESSAGES', 'messages.po')
            if not self.pod.file_exists(po_path):
                skipped['po'].add(str(locale))
                continue

            clean_locales.append(locale)

        # Summary of skipped files.
        if skipped['symlink']:
            self.pod.logger.info('Skipping: {} (symlinked)'.format(
                ', '.join(sorted(skipped['symlink']))))
        if skipped['po']:
            self.pod.logger.info('Skipping: {} (no `.po` file)'.format(
                ', '.join(sorted(skipped['po']))))

        return clean_locales

    def _download_content(self, stat):
        raise NotImplementedError

    def _upload_catalog(self, catalog, source_lang, prune):
        raise NotImplementedError

    def _upload_catalogs(self, catalogs, source_lang, prune=False):
        raise NotImplementedError

    def _update_acl(self, stat, locale):
        raise NotImplementedError

    def _update_acls(self, stat, locales):
        raise NotImplementedError

    def _update_meta(self, stat, locale, catalog):
        raise NotImplementedError

    def _get_stats_to_download(self, locales):
        # 'stats' maps the service name to a mapping of languages to stats.
        if not self.pod.file_exists(Translator.TRANSLATOR_STATS_PATH):
            return {}
        stats = self.pod.read_yaml(Translator.TRANSLATOR_STATS_PATH)
        if self.KIND not in stats:
            self.pod.logger.info(
                'Nothing found to download from {}'.format(self.KIND))
            return {}
        stats_to_download = stats[self.KIND]
        if locales:
            stats_to_download = dict([(lang, stat)
                                      for (lang, stat) in stats_to_download.iteritems()
                                      if lang in locales])
        for lang, stat in stats_to_download.iteritems():
            stat['lang'] = lang
            stat = json.dumps(stat)
            stat_message = protojson.decode_message(TranslatorStat, stat)
            stats_to_download[lang] = stat_message
        return stats_to_download

    def download(self, locales, save_stats=True, inject=False, include_obsolete=False):
        # TODO: Rename to `download_and_import`.
        if not self.pod.file_exists(Translator.TRANSLATOR_STATS_PATH):
            text = 'File {} not found. Nothing to download.'
            self.pod.logger.info(text.format(Translator.TRANSLATOR_STATS_PATH))
            return
        stats_to_download = self._get_stats_to_download(locales)
        if not stats_to_download:
            return
        num_files = len(stats_to_download)
        text = 'Downloading translations: %(value)d/{} (in %(time_elapsed).9s)'
        widgets = [progressbar.FormatLabel(text.format(num_files))]
        if not inject:
            bar = progressbar_non.create_progressbar(
                "Downloading translations...", widgets=widgets, max_value=num_files)
            bar.start()
        threads = []
        langs_to_translations = {}
        new_stats = []

        def _do_download(lang, stat):
            try:
                new_stat, content = self._download_content(stat)
            except translator_errors.NotFoundError:
                text = 'No translations to download for: {}'
                self.pod.logger.info(text.format(lang))
                return

            new_stat.uploaded = stat.uploaded  # Preserve uploaded field.
            langs_to_translations[lang] = content
            new_stats.append(new_stat)
        for i, (lang, stat) in enumerate(stats_to_download.iteritems()):
            if inject:
                thread = threading.Thread(
                    target=_do_download, args=(lang, stat))
            else:
                thread = utils.ProgressBarThread(
                    bar, True, target=_do_download, args=(lang, stat))
            threads.append(thread)
            thread.start()
            # Perform the first operation synchronously to avoid oauth2 refresh
            # locking issues.
            if i == 0:
                thread.join()
        for i, thread in enumerate(threads):
            if i > 0:
                thread.join()
        if not inject:
            bar.finish()

        has_changed_content = False
        for lang, translations in langs_to_translations.iteritems():
            if inject:
                if self.pod.catalogs.inject_translations(locale=lang, content=translations):
                    has_changed_content = True
            elif self.pod.catalogs.import_translations(
                    locale=lang, content=translations,
                    include_obsolete=include_obsolete):
                has_changed_content = True

        if save_stats and has_changed_content:
            self.save_stats(new_stats)
        return new_stats

    def update_acl(self, locales=None):
        locales = locales or self.pod.catalogs.list_locales()
        locales = self._cleanup_locales(locales)
        if not locales:
            self.pod.logger.info('No locales to found to update.')
            return
        stats_to_download = self._get_stats_to_download(locales)
        if not stats_to_download:
            self.pod.logger.info('No documents found to update.')
            return
        if self.has_multiple_langs_in_one_resource:
            self._update_acls(stats_to_download, locales)
            stat = stats_to_download.values()[0]
            self.pod.logger.info('ACL updated -> {}'.format(stat.ident))
            return
        threads = []
        for i, (locale, stat) in enumerate(stats_to_download.iteritems()):
            thread = threading.Thread(
                target=self._update_acl, args=(stat, locale))
            threads.append(thread)
            thread.start()
            if i == 0:
                thread.join()
            self.pod.logger.info(
                'ACL updated ({}): {}'.format(stat.lang, stat.url))
        for i, thread in enumerate(threads):
            if i > 0:
                thread.join()

    def update_meta(self, locales=None):
        locales = locales or self.pod.catalogs.list_locales()
        locales = self._cleanup_locales(locales)
        if not locales:
            self.pod.logger.info('No locales to found to update.')
            return
        stats_to_download = self._get_stats_to_download(locales)
        if not stats_to_download:
            self.pod.logger.info('No documents found to update.')
            return
        threads = []
        for i, (locale, stat) in enumerate(stats_to_download.iteritems()):
            catalog_for_meta = self.pod.catalogs.get(locale)
            thread = threading.Thread(
                target=self._update_meta, args=(stat, locale, catalog_for_meta))
            threads.append(thread)
            thread.start()
            if i == 0:
                thread.join()
            self.pod.logger.info('Meta information updated ({}): {}'.format(
                stat.lang, stat.url))
        for i, thread in enumerate(threads):
            if i > 0:
                thread.join()

    def upload(self, locales=None, force=True, verbose=False, save_stats=True,
               prune=False):
        source_lang = self.pod.podspec.default_locale
        locales = locales or self.pod.catalogs.list_locales()
        locales = self._cleanup_locales(locales)
        stats = []
        num_files = len(locales)
        if not locales:
            self.pod.logger.info('No locales to upload.')
            return
        if not force:
            if (self.has_immutable_translation_resources
                    and self.pod.file_exists(Translator.TRANSLATOR_STATS_PATH)):
                text = 'Found existing translator data in: {}'
                self.pod.logger.info(text.format(
                    Translator.TRANSLATOR_STATS_PATH))
                text = 'This will be updated with new data after the upload is complete.'
                self.pod.logger.info(text)
            text = 'Proceed to upload {} translation catalogs?'
            text = text.format(num_files)
            if not utils.interactive_confirm(text):
                self.pod.logger.info('Aborted.')
                return
        if self.has_multiple_langs_in_one_resource:
            catalogs_to_upload = []
            for locale in locales:
                catalog_to_upload = self.pod.catalogs.get(locale)
                if catalog_to_upload:
                    catalogs_to_upload.append(catalog_to_upload)
            stats = self._upload_catalogs(catalogs_to_upload, source_lang,
                                          prune=prune)
        else:
            text = 'Uploading translations: %(value)d/{} (in %(time_elapsed).9s)'
            widgets = [progressbar.FormatLabel(text.format(num_files))]
            bar = progressbar_non.create_progressbar(
                "Uploading translations...", widgets=widgets, max_value=num_files)
            bar.start()
            threads = []

            def _do_upload(locale):
                catalog = self.pod.catalogs.get(locale)
                stat = self._upload_catalog(catalog, source_lang, prune=prune)
                stats.append(stat)
            for i, locale in enumerate(locales):
                thread = utils.ProgressBarThread(
                    bar, True, target=_do_upload, args=(locale,))
                threads.append(thread)
                thread.start()
                # Perform the first operation synchronously to avoid oauth2 refresh
                # locking issues.
                if i == 0:
                    thread.join()
            for i, thread in enumerate(threads):
                if i > 0:
                    thread.join()
            bar.finish()
        stats = sorted(stats, key=lambda stat: stat.lang)
        if verbose:
            self.pretty_print_stats(stats)
        if save_stats:
            self.save_stats(stats)
        return stats

    def save_stats(self, stats):
        """Merges a list of stats into the translator stats file."""
        if self.pod.file_exists(Translator.TRANSLATOR_STATS_PATH):
            content = self.pod.read_yaml(Translator.TRANSLATOR_STATS_PATH)
            create = False
        else:
            content = {}
            create = True
        if self.KIND not in content:
            content[self.KIND] = {}
        for stat in copy.deepcopy(stats):
            content[self.KIND][stat.lang] = stat
        yaml_content = yaml.safe_dump(content, default_flow_style=False)
        self.pod.write_file(Translator.TRANSLATOR_STATS_PATH, yaml_content)
        if create:
            self.pod.logger.info('Saved: {}'.format(
                Translator.TRANSLATOR_STATS_PATH))
        else:
            self.pod.logger.info('Updated: {}'.format(
                Translator.TRANSLATOR_STATS_PATH))

    @classmethod
    def pretty_print_stats(cls, stats):
        table = texttable.Texttable(max_width=0)
        table.set_deco(texttable.Texttable.HEADER)
        rows = []
        rows.append(['Language', 'URL', 'Wordcount'])
        for stat in stats:
            rows.append([stat.lang, stat.url, stat.num_words or '--'])
        table.add_rows(rows)
        logging.info('\n' + table.draw() + '\n')

    def get_edit_url(self, doc):
        if not doc.locale:
            return
        stats = self._get_stats_to_download([doc.locale])
        if doc.locale not in stats:
            return
        stat = stats[doc.locale]
        return stat.url

    def inject(self, doc):
        if not self._inject or not doc.locale:
            return
        self.download(locales=[doc.locale], save_stats=False, inject=True)
        return self
<commit_msg>Fixing the case when the stat is already a TranslatorStat.<commit_after>"""Base translators for translating Grow content."""

import copy
import json
import logging
import os
import threading
import progressbar
import texttable
import yaml
from protorpc import message_types
from protorpc import messages
from protorpc import protojson
from grow.common import progressbar_non
from grow.common import utils
from grow.translators import errors as translator_errors

class TranslatorStat(messages.Message):
    lang = messages.StringField(1)
    num_words = messages.IntegerField(2)
    num_words_translated = messages.IntegerField(3)
    source_lang = messages.StringField(4)
    ident = messages.StringField(5)
    url = messages.StringField(6)
    uploaded = message_types.DateTimeField(7)
    service = messages.StringField(8)
    downloaded = message_types.DateTimeField(9)

def translator_stat_representer(dumper, stat):
    content = json.loads(protojson.encode_message(stat))
    content.pop('lang')  # Exclude from serialization.
    return dumper.represent_mapping('tag:yaml.org,2002:map', content)

yaml.SafeDumper.add_representer(TranslatorStat,
                                translator_stat_representer)

class TranslatorServiceError(Exception):

    def __init__(self, message, ident=None, locale=None):
        if locale:
            new_message = 'Error for locale "{}" -> {}'.format(locale, message)
        elif ident:
            new_message = 'Error for resource "{}" -> {}'.format(
                ident, message)
        else:
            new_message = message
        super(TranslatorServiceError, self).__init__(new_message)

class Translator(object):
    TRANSLATOR_STATS_PATH = '/translators.yaml'
    KIND = None
    has_immutable_translation_resources = False
    has_multiple_langs_in_one_resource = False

    def __init__(self, pod, config=None, project_title=None,
                 instructions=None, inject=False):
        self.pod = pod
        self.config = config or {}
        self.project_title = project_title or 'Untitled Grow Website'
        self.instructions = instructions
        self._inject = inject

    def _cleanup_locales(self, locales):
        """Certain locales should be ignored."""
        clean_locales = []
        default_locale = self.pod.podspec.default_locale
        skipped = {
            'symlink': set(),
            'po': set(),
        }
        for locale in locales:
            locale_path = os.path.join('translations', str(locale))

            # Silently ignore the default locale.
            if default_locale and str(locale) == str(default_locale):
                continue

            # Ignore the symlinks.
            if os.path.islink(locale_path):
                skipped['symlink'].add(str(locale))
                continue

            # Ignore the locales without a `.PO` file.
            po_path = os.path.join(locale_path, 'LC_MESSAGES', 'messages.po')
            if not self.pod.file_exists(po_path):
                skipped['po'].add(str(locale))
                continue

            clean_locales.append(locale)

        # Summary of skipped files.
        if skipped['symlink']:
            self.pod.logger.info('Skipping: {} (symlinked)'.format(
                ', '.join(sorted(skipped['symlink']))))
        if skipped['po']:
            self.pod.logger.info('Skipping: {} (no `.po` file)'.format(
                ', '.join(sorted(skipped['po']))))

        return clean_locales

    def _download_content(self, stat):
        raise NotImplementedError

    def _upload_catalog(self, catalog, source_lang, prune):
        raise NotImplementedError

    def _upload_catalogs(self, catalogs, source_lang, prune=False):
        raise NotImplementedError

    def _update_acl(self, stat, locale):
        raise NotImplementedError

    def _update_acls(self, stat, locales):
        raise NotImplementedError

    def _update_meta(self, stat, locale, catalog):
        raise NotImplementedError

    def _get_stats_to_download(self, locales):
        # 'stats' maps the service name to a mapping of languages to stats.
        if not self.pod.file_exists(Translator.TRANSLATOR_STATS_PATH):
            return {}
        stats = self.pod.read_yaml(Translator.TRANSLATOR_STATS_PATH)
        if self.KIND not in stats:
            self.pod.logger.info(
                'Nothing found to download from {}'.format(self.KIND))
            return {}
        stats_to_download = stats[self.KIND]
        if locales:
            stats_to_download = dict([(lang, stat)
                                      for (lang, stat) in stats_to_download.iteritems()
                                      if lang in locales])
        for lang, stat in stats_to_download.iteritems():
            if isinstance(stat, TranslatorStat):
                stat = json.loads(protojson.encode_message(stat))
            stat['lang'] = lang
            stat_message = protojson.decode_message(TranslatorStat, json.dumps(stat))
            stats_to_download[lang] = stat_message
        return stats_to_download

    def download(self, locales, save_stats=True, inject=False, include_obsolete=False):
        # TODO: Rename to `download_and_import`.
        if not self.pod.file_exists(Translator.TRANSLATOR_STATS_PATH):
            text = 'File {} not found. Nothing to download.'
            self.pod.logger.info(text.format(Translator.TRANSLATOR_STATS_PATH))
            return
        stats_to_download = self._get_stats_to_download(locales)
        if not stats_to_download:
            return
        num_files = len(stats_to_download)
        text = 'Downloading translations: %(value)d/{} (in %(time_elapsed).9s)'
        widgets = [progressbar.FormatLabel(text.format(num_files))]
        if not inject:
            bar = progressbar_non.create_progressbar(
                "Downloading translations...", widgets=widgets, max_value=num_files)
            bar.start()
        threads = []
        langs_to_translations = {}
        new_stats = []

        def _do_download(lang, stat):
            try:
                new_stat, content = self._download_content(stat)
            except translator_errors.NotFoundError:
                text = 'No translations to download for: {}'
                self.pod.logger.info(text.format(lang))
                return

            new_stat.uploaded = stat.uploaded  # Preserve uploaded field.
            langs_to_translations[lang] = content
            new_stats.append(new_stat)
        for i, (lang, stat) in enumerate(stats_to_download.iteritems()):
            if inject:
                thread = threading.Thread(
                    target=_do_download, args=(lang, stat))
            else:
                thread = utils.ProgressBarThread(
                    bar, True, target=_do_download, args=(lang, stat))
            threads.append(thread)
            thread.start()
            # Perform the first operation synchronously to avoid oauth2 refresh
            # locking issues.
            if i == 0:
                thread.join()
        for i, thread in enumerate(threads):
            if i > 0:
                thread.join()
        if not inject:
            bar.finish()

        has_changed_content = False
        for lang, translations in langs_to_translations.iteritems():
            if inject:
                if self.pod.catalogs.inject_translations(locale=lang, content=translations):
                    has_changed_content = True
            elif self.pod.catalogs.import_translations(
                    locale=lang, content=translations,
                    include_obsolete=include_obsolete):
                has_changed_content = True

        if save_stats and has_changed_content:
            self.save_stats(new_stats)
        return new_stats

    def update_acl(self, locales=None):
        locales = locales or self.pod.catalogs.list_locales()
        locales = self._cleanup_locales(locales)
        if not locales:
            self.pod.logger.info('No locales to found to update.')
            return
        stats_to_download = self._get_stats_to_download(locales)
        if not stats_to_download:
            self.pod.logger.info('No documents found to update.')
            return
        if self.has_multiple_langs_in_one_resource:
            self._update_acls(stats_to_download, locales)
            stat = stats_to_download.values()[0]
            self.pod.logger.info('ACL updated -> {}'.format(stat.ident))
            return
        threads = []
        for i, (locale, stat) in enumerate(stats_to_download.iteritems()):
            thread = threading.Thread(
                target=self._update_acl, args=(stat, locale))
            threads.append(thread)
            thread.start()
            if i == 0:
                thread.join()
            self.pod.logger.info(
                'ACL updated ({}): {}'.format(stat.lang, stat.url))
        for i, thread in enumerate(threads):
            if i > 0:
                thread.join()

    def update_meta(self, locales=None):
        locales = locales or self.pod.catalogs.list_locales()
        locales = self._cleanup_locales(locales)
        if not locales:
            self.pod.logger.info('No locales to found to update.')
            return
        stats_to_download = self._get_stats_to_download(locales)
        if not stats_to_download:
            self.pod.logger.info('No documents found to update.')
            return
        threads = []
        for i, (locale, stat) in enumerate(stats_to_download.iteritems()):
            catalog_for_meta = self.pod.catalogs.get(locale)
            thread = threading.Thread(
                target=self._update_meta, args=(stat, locale, catalog_for_meta))
            threads.append(thread)
            thread.start()
            if i == 0:
                thread.join()
            self.pod.logger.info('Meta information updated ({}): {}'.format(
                stat.lang, stat.url))
        for i, thread in enumerate(threads):
            if i > 0:
                thread.join()

    def upload(self, locales=None, force=True, verbose=False, save_stats=True,
               prune=False):
        source_lang = self.pod.podspec.default_locale
        locales = locales or self.pod.catalogs.list_locales()
        locales = self._cleanup_locales(locales)
        stats = []
        num_files = len(locales)
        if not locales:
            self.pod.logger.info('No locales to upload.')
            return
        if not force:
            if (self.has_immutable_translation_resources
                    and self.pod.file_exists(Translator.TRANSLATOR_STATS_PATH)):
                text = 'Found existing translator data in: {}'
                self.pod.logger.info(text.format(
                    Translator.TRANSLATOR_STATS_PATH))
                text = 'This will be updated with new data after the upload is complete.'
                self.pod.logger.info(text)
            text = 'Proceed to upload {} translation catalogs?'
            text = text.format(num_files)
            if not utils.interactive_confirm(text):
                self.pod.logger.info('Aborted.')
                return
        if self.has_multiple_langs_in_one_resource:
            catalogs_to_upload = []
            for locale in locales:
                catalog_to_upload = self.pod.catalogs.get(locale)
                if catalog_to_upload:
                    catalogs_to_upload.append(catalog_to_upload)
            stats = self._upload_catalogs(catalogs_to_upload, source_lang,
                                          prune=prune)
        else:
            text = 'Uploading translations: %(value)d/{} (in %(time_elapsed).9s)'
            widgets = [progressbar.FormatLabel(text.format(num_files))]
            bar = progressbar_non.create_progressbar(
                "Uploading translations...", widgets=widgets, max_value=num_files)
            bar.start()
            threads = []

            def _do_upload(locale):
                catalog = self.pod.catalogs.get(locale)
                stat = self._upload_catalog(catalog, source_lang, prune=prune)
                stats.append(stat)
            for i, locale in enumerate(locales):
                thread = utils.ProgressBarThread(
                    bar, True, target=_do_upload, args=(locale,))
                threads.append(thread)
                thread.start()
                # Perform the first operation synchronously to avoid oauth2 refresh
                # locking issues.
                if i == 0:
                    thread.join()
            for i, thread in enumerate(threads):
                if i > 0:
                    thread.join()
            bar.finish()
        stats = sorted(stats, key=lambda stat: stat.lang)
        if verbose:
            self.pretty_print_stats(stats)
        if save_stats:
            self.save_stats(stats)
        return stats

    def save_stats(self, stats):
        """Merges a list of stats into the translator stats file."""
        if self.pod.file_exists(Translator.TRANSLATOR_STATS_PATH):
            content = self.pod.read_yaml(Translator.TRANSLATOR_STATS_PATH)
            create = False
        else:
            content = {}
            create = True
        if self.KIND not in content:
            content[self.KIND] = {}
        for stat in copy.deepcopy(stats):
            content[self.KIND][stat.lang] = stat
        yaml_content = yaml.safe_dump(content, default_flow_style=False)
        self.pod.write_file(Translator.TRANSLATOR_STATS_PATH, yaml_content)
        if create:
            self.pod.logger.info('Saved: {}'.format(
                Translator.TRANSLATOR_STATS_PATH))
        else:
            self.pod.logger.info('Updated: {}'.format(
                Translator.TRANSLATOR_STATS_PATH))

    @classmethod
    def pretty_print_stats(cls, stats):
        table = texttable.Texttable(max_width=0)
        table.set_deco(texttable.Texttable.HEADER)
        rows = []
        rows.append(['Language', 'URL', 'Wordcount'])
        for stat in stats:
            rows.append([stat.lang, stat.url, stat.num_words or '--'])
        table.add_rows(rows)
        logging.info('\n' + table.draw() + '\n')

    def get_edit_url(self, doc):
        if not doc.locale:
            return
        stats = self._get_stats_to_download([doc.locale])
        if doc.locale not in stats:
            return
        stat = stats[doc.locale]
        return stat.url

    def inject(self, doc):
        if not self._inject or not doc.locale:
            return
        self.download(locales=[doc.locale], save_stats=False, inject=True)
        return self
<|endoftext|>


<commit_before>from django.db import models

from filingcabinet.models import (
    AbstractDocument,
    AbstractDocumentCollection,
    get_page_image_filename
)

from froide.helper.auth import (
    can_read_object, can_write_object, get_write_queryset
)

class Document(AbstractDocument):
    original = models.ForeignKey(
        'foirequest.FoiAttachment', null=True, blank=True,
        on_delete=models.SET_NULL, related_name='original_document'
    )

    foirequest = models.ForeignKey(
        'foirequest.FoiRequest', null=True, blank=True,
        on_delete=models.SET_NULL
    )
    publicbody = models.ForeignKey(
        'publicbody.PublicBody', null=True, blank=True,
        on_delete=models.SET_NULL
    )
    team = models.ForeignKey(
        'team.Team', null=True, blank=True,
        on_delete=models.SET_NULL
    )

    def is_public(self):
        return self.public

    def can_read(self, request):
        return can_read_object(self, request=request)

    def can_write(self, request):
        return can_write_object(self, request=request)

    @classmethod
    def get_annotatable(cls, request):
        cond = models.Q(public=True, allow_annotation=True)
        return get_write_queryset(
            cls.objects.all(),
            request,
            has_team=True,
            user_write_filter=cond,
        )

    def get_serializer_class(self, detail=False):
        from .api_views import DocumentSerializer, DocumentDetailSerializer

        if detail:
            return DocumentDetailSerializer
        return DocumentSerializer

    def get_crossdomain_auth(self, filename=None):
        from .auth import DocumentCrossDomainMediaAuth

        if filename is None:
            filename = self.get_document_filename()

        return DocumentCrossDomainMediaAuth({
            'object': self,
            'filename': filename
        })

    def get_authorized_file_url(self, filename=None):
        if self.public:
            return self.get_file_url(filename=filename)
        return self.get_crossdomain_auth(filename=filename).get_full_media_url(
            authorized=True
        )

    def get_page_template(self, page=1, size='small'):
        filename = get_page_image_filename(page=page, size=size)
        return self.get_authorized_file_url(filename=filename)

    def get_cover_image(self):
        return self.get_authorized_file_url(filename=get_page_image_filename(
            page=1, size='small'
        ))

class DocumentCollection(AbstractDocumentCollection):
    team = models.ForeignKey(
        'team.Team', null=True, blank=True,
        on_delete=models.SET_NULL
    )

    def is_public(self):
        return self.public

    def get_serializer_class(self):
        from .api_views import DocumentCollectionSerializer

        return DocumentCollectionSerializer
<commit_msg>Fix bad size parameter on document page template<commit_after>from django.db import models

from filingcabinet.models import (
    AbstractDocument,
    AbstractDocumentCollection,
    get_page_image_filename
)

from froide.helper.auth import (
    can_read_object, can_write_object, get_write_queryset
)

class Document(AbstractDocument):
    original = models.ForeignKey(
        'foirequest.FoiAttachment', null=True, blank=True,
        on_delete=models.SET_NULL, related_name='original_document'
    )

    foirequest = models.ForeignKey(
        'foirequest.FoiRequest', null=True, blank=True,
        on_delete=models.SET_NULL
    )
    publicbody = models.ForeignKey(
        'publicbody.PublicBody', null=True, blank=True,
        on_delete=models.SET_NULL
    )
    team = models.ForeignKey(
        'team.Team', null=True, blank=True,
        on_delete=models.SET_NULL
    )

    def is_public(self):
        return self.public

    def can_read(self, request):
        return can_read_object(self, request=request)

    def can_write(self, request):
        return can_write_object(self, request=request)

    @classmethod
    def get_annotatable(cls, request):
        cond = models.Q(public=True, allow_annotation=True)
        return get_write_queryset(
            cls.objects.all(),
            request,
            has_team=True,
            user_write_filter=cond,
        )

    def get_serializer_class(self, detail=False):
        from .api_views import DocumentSerializer, DocumentDetailSerializer

        if detail:
            return DocumentDetailSerializer
        return DocumentSerializer

    def get_crossdomain_auth(self, filename=None):
        from .auth import DocumentCrossDomainMediaAuth

        if filename is None:
            filename = self.get_document_filename()

        return DocumentCrossDomainMediaAuth({
            'object': self,
            'filename': filename
        })

    def get_authorized_file_url(self, filename=None):
        if self.public:
            return self.get_file_url(filename=filename)
        return self.get_crossdomain_auth(filename=filename).get_full_media_url(
            authorized=True
        )

    def get_page_template(self, page=1, size='{size}'):
        filename = get_page_image_filename(page=page, size=size)
        return self.get_authorized_file_url(filename=filename)

    def get_cover_image(self):
        return self.get_authorized_file_url(filename=get_page_image_filename(
            page=1, size='small'
        ))

class DocumentCollection(AbstractDocumentCollection):
    team = models.ForeignKey(
        'team.Team', null=True, blank=True,
        on_delete=models.SET_NULL
    )

    def is_public(self):
        return self.public

    def get_serializer_class(self):
        from .api_views import DocumentCollectionSerializer

        return DocumentCollectionSerializer
<|endoftext|>


<commit_before>#!/usr/bin/env python
#
# Copyright 2012 Google Inc.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Marc Cohen
#

'''Render Home Page.'''

import json
import logging
import pickle
from google.appengine.api import memcache
from google.appengine.api import users
from google.appengine.ext import webapp
from google.appengine.ext.webapp import template
from google.appengine.ext.webapp.util import login_required
from oauth2client.appengine import StorageByKeyName
from oauth2client.appengine import CredentialsModel
from oauth2client.client import OAuth2WebServerFlow

# NOTE: You need to edit the redirect URI into the Dev Console, and that value
# needs to be updated when you move from running this app locally to running
# on Google App Engine (else you'll get a 'redirect_uri_mismatch' error).

# Static configuation constants.
SCOPE = ('https://www.googleapis.com/auth/devstorage.read_write ' +
         'https://www.googleapis.com/auth/prediction')
USER_AGENT = 'try-prediction/1.0'
DEFAULT_MODEL = 'Language Detection'
SECRETS_FILE = 'rc/client_secrets.json'
MODELS_FILE = 'rc/models.json'

def parse_json_file(file):
  '''Utility function to open, read, and parse the contents of
     a file containing text encoded as a JSON document, and
     return resulting json object to caller.
  '''
  f = open(file, 'r')
  json_str = f.read()
  f.close()
  return json.loads(json_str)

class HomePage(webapp.RequestHandler):
  '''This class renders the main home page for the "Try Prediction" app.'''

  def post(self):
    '''Use the same logic for posts and gets.'''
    self.get()

  def get(self):
    '''Process get requests.'''
    user = users.get_current_user()

    # Fetch stored server credentials.
    credentials = StorageByKeyName(CredentialsModel, USER_AGENT,
                                   'credentials').locked_get()

    # If server credentials not found, trigger OAuth2.0 web server flow.
    if not credentials or credentials.invalid:
      # If no credentials and no user logged in, redirect to the reset,
      # which will force a login to make sure this user has permission
      # to initialize the shared server credentials.
      if not user:
        self.redirect("/reset")

      # Read and parse client secrets JSON file.
      secrets = parse_json_file(SECRETS_FILE)

      client_id = secrets['installed']['client_id']
      client_secret = secrets['installed']['client_secret']

      flow = OAuth2WebServerFlow(client_id=client_id,
                                 client_secret=client_secret,
                                 scope=SCOPE,
                                 user_agent=USER_AGENT,
                                 access_type = 'offline',
                                 approval_prompt='force')
      callback = self.request.relative_url('/auth_return')
      authorize_url = flow.step1_get_authorize_url(callback)

      # Save flow object in memcache for later retrieval on OAuth callback,
      # and redirect this session to Google's OAuth 2.0 authorization site.
      logging.info('saving flow for user ' + user.user_id())
      memcache.set(user.user_id(), pickle.dumps(flow))
      self.redirect(authorize_url)

    # At this point we should have valid server credentials stored away.

    # Get user's chosen model, if selected, or default model.
    selected_model = self.request.get('model', DEFAULT_MODEL)

    # Read and parse model description JSON file.
    models = parse_json_file(MODELS_FILE)

    # Set django template and render home page.
    template_values = {
      'user' : str(user),
      'selected_model' : selected_model,
      'models' : models,
    }
    self.response.out.write(template.render('home.html', template_values))

class AuthHandler(webapp.RequestHandler):
  '''This class fields OAuth 2.0 web callback for the "Try Prediction" app.'''

  @login_required
  def get(self):
    user = users.get_current_user()

    # Retrieve flow object from memcache.
    logging.info('retrieving flow for user ' + user.user_id())
    flow = pickle.loads(memcache.get(user.user_id()))
    if flow:
      # Extract newly acquired server credentials, store creds
      # in datatore for future retrieval and redirect session
      # back to app's main page.
      credentials = flow.step2_exchange(self.request.params)
      StorageByKeyName(CredentialsModel, USER_AGENT,
                       'credentials').locked_put(credentials)
      self.redirect('/')
    else:
      raise('unable to obtain OAuth 2.0 credentials')

class Reset(webapp.RequestHandler):
  '''This class processes requests to reset the server's OAuth 2.0
     credentials. It should only be executed by the application
     administrator per the app.yaml configuration file.'''

  @login_required
  def get(self):
    # Store empty credentials in the datastore and redirect to main page.
    StorageByKeyName(CredentialsModel, USER_AGENT,
                     'credentials').locked_put(None)
    self.redirect('/')
<commit_msg>Abort imidiatelly and redirect.<commit_after>#!/usr/bin/env python
#
# Copyright 2012 Google Inc.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Marc Cohen
#

'''Render Home Page.'''

import json
import logging
import pickle
from google.appengine.api import memcache
from google.appengine.api import users
from google.appengine.ext import webapp
from google.appengine.ext.webapp import template
from google.appengine.ext.webapp.util import login_required
from oauth2client.appengine import StorageByKeyName
from oauth2client.appengine import CredentialsModel
from oauth2client.client import OAuth2WebServerFlow

# NOTE: You need to edit the redirect URI into the Dev Console, and that value
# needs to be updated when you move from running this app locally to running
# on Google App Engine (else you'll get a 'redirect_uri_mismatch' error).

# Static configuation constants.
SCOPE = ('https://www.googleapis.com/auth/devstorage.read_write ' +
         'https://www.googleapis.com/auth/prediction')
USER_AGENT = 'try-prediction/1.0'
DEFAULT_MODEL = 'Language Detection'
SECRETS_FILE = 'rc/client_secrets.json'
MODELS_FILE = 'rc/models.json'

def parse_json_file(file):
  '''Utility function to open, read, and parse the contents of
     a file containing text encoded as a JSON document, and
     return resulting json object to caller.
  '''
  f = open(file, 'r')
  json_str = f.read()
  f.close()
  return json.loads(json_str)

class HomePage(webapp.RequestHandler):
  '''This class renders the main home page for the "Try Prediction" app.'''

  def post(self):
    '''Use the same logic for posts and gets.'''
    self.get()

  def get(self):
    '''Process get requests.'''
    user = users.get_current_user()

    # Fetch stored server credentials.
    credentials = StorageByKeyName(CredentialsModel, USER_AGENT,
                                   'credentials').locked_get()

    # If server credentials not found, trigger OAuth2.0 web server flow.
    if not credentials or credentials.invalid:
      # If no credentials and no user logged in, redirect to the reset,
      # which will force a login to make sure this user has permission
      # to initialize the shared server credentials.
      if not user:
        self.redirect("/reset", abort=True)

      # Read and parse client secrets JSON file.
      secrets = parse_json_file(SECRETS_FILE)

      client_id = secrets['installed']['client_id']
      client_secret = secrets['installed']['client_secret']

      flow = OAuth2WebServerFlow(client_id=client_id,
                                 client_secret=client_secret,
                                 scope=SCOPE,
                                 user_agent=USER_AGENT,
                                 access_type = 'offline',
                                 approval_prompt='force')
      callback = self.request.relative_url('/auth_return')
      authorize_url = flow.step1_get_authorize_url(callback)

      # Save flow object in memcache for later retrieval on OAuth callback,
      # and redirect this session to Google's OAuth 2.0 authorization site.
      logging.info('saving flow for user ' + user.user_id())
      memcache.set(user.user_id(), pickle.dumps(flow))
      self.redirect(authorize_url)

    # At this point we should have valid server credentials stored away.

    # Get user's chosen model, if selected, or default model.
    selected_model = self.request.get('model', DEFAULT_MODEL)

    # Read and parse model description JSON file.
    models = parse_json_file(MODELS_FILE)

    # Set django template and render home page.
    template_values = {
      'user' : str(user),
      'selected_model' : selected_model,
      'models' : models,
    }
    self.response.out.write(template.render('home.html', template_values))

class AuthHandler(webapp.RequestHandler):
  '''This class fields OAuth 2.0 web callback for the "Try Prediction" app.'''

  @login_required
  def get(self):
    user = users.get_current_user()

    # Retrieve flow object from memcache.
    logging.info('retrieving flow for user ' + user.user_id())
    flow = pickle.loads(memcache.get(user.user_id()))
    if flow:
      # Extract newly acquired server credentials, store creds
      # in datatore for future retrieval and redirect session
      # back to app's main page.
      credentials = flow.step2_exchange(self.request.params)
      StorageByKeyName(CredentialsModel, USER_AGENT,
                       'credentials').locked_put(credentials)
      self.redirect('/')
    else:
      raise('unable to obtain OAuth 2.0 credentials')

class Reset(webapp.RequestHandler):
  '''This class processes requests to reset the server's OAuth 2.0
     credentials. It should only be executed by the application
     administrator per the app.yaml configuration file.'''

  @login_required
  def get(self):
    # Store empty credentials in the datastore and redirect to main page.
    StorageByKeyName(CredentialsModel, USER_AGENT,
                     'credentials').locked_put(None)
    self.redirect('/')
<|endoftext|>


<commit_before>#!/usr/bin/env ccp4-python

'''
09.11.2015

@author: hlfsimko
'''

# System
import argparse
import collections
import logging
import os
import pickle
import sys
import tempfile
import unittest

# 3rd party
import numpy

if not "CCP4" in os.environ.keys(): raise RuntimeError('CCP4 not found')
sys.path.insert(0, os.path.join(os.environ['CCP4'], "share", "ample", "parsers"))
#sys.path.insert(0, os.path.join(os.environ["HOME"], "opt", "ample-dev1", "parsers"))

# Custom
import ample_util
import parse_tmscore
import pdb_edit

try:
    import parse_alignment
    _BIOPYTHON = True
except ImportError:
    _BIOPYTHON = False

class TMscorer(object):
    def __init__(self, structure, tmscore_exe, wdir=None):
        self.logger = logging.getLogger()
        self.pickle_file = None
        self.structure = structure
        self.tmscore_exe = tmscore_exe
        self.workingDIR = wdir
        return

    def main(self, pdb_list_file):
        # Read all pdb files from list
        pdbs = self._read_list(pdb_list_file)

        # Compare each pdb in list against structure defined
        self.entries = self.compare_to_structure(pdbs)

        # Dumpe all data in a pickle file
        self.pickle_file = os.path.join(self.workingDIR, "tmresults.pkl")
        pickle.dump(self.entries, open(self.pickle_file, 'w'))
        return

    def compare_to_structure(self, pdb_list):
        self.logger.info('-------Evaluating decoys/models-------')

        entries = []

        for pdb in pdb_list:
            if not os.path.exists(pdb): continue

            name = os.path.basename(pdb).rsplit(".", 1)[0]
            pdbin = os.path.abspath(pdb)

            self.logger.info("Working on %s" % name)

            # Write a temporary file with the aligned residues from the native structure
            structure_mod = tempfile.NamedTemporaryFile(delete=False)
            structure_mod.close()
            pdbin_mod = tempfile.NamedTemporaryFile(delete=False)
            pdbin_mod.close()

            # Initialise the logparser that stores all the scores
            pt = parse_tmscore.TMscoreLogParser()

            # Do the try clause here to allow anything that is required from here to throw
            # exceptions. In that case we revert to the TMscoreLogParser default values of
            # 0.0 for every score.
            try:
                self.mod_structures(pdbin, pdbin_mod.name, self.structure, structure_mod.name)

                # Create a command list and execute TMscore
                log = os.path.join(self.workingDIR, name+".tmscore.log")
                cmd = [ self.tmscore_exe, pdbin_mod.name, structure_mod.name ]
                p = ample_util.run_command(cmd, logfile=log, directory=self.workingDIR)

                os.unlink(structure_mod.name)
                os.unlink(pdbin_mod.name)

                # Parse the log file to the parser
                pt.parse(log)

            except Exception:
                log = "None"

            entry = self._store(name, pdbin, log, self.structure, pt)
            entries.append(entry)

        return entries

    def mod_structures(self, pdbin, pdbin_mod, structure, structure_mod):
        ''' Make sure the decoy and the xtal pdb align to get an accurate Tm score '''

        if not _BIOPYTHON:
            raise ImportError

        # Disable the info logger to not spam the user with which chain of native extracted.
        # Happens for every model + native below
        # http://stackoverflow.com/questions/2266646/how-to-i-disable-and-re-enable-console-logging-in-python
        logging.disable( logging.INFO )

        pdbin_seq     = pdb_edit.sequence(pdbin).values()[0]
        structure_seq = pdb_edit.sequence(structure).values()[0]

        # Align the sequences to see how much of the predicted decoys are in the xtal
        aligned_seq_list = parse_alignment.AlignmentParser().align_sequences(pdbin_seq, structure_seq)
        pdbin_seq_ali     = aligned_seq_list[0]
        structure_seq_ali = aligned_seq_list[1]

        # Get the gaps in both sequences
        pdbin_gaps     = self.find_gaps(pdbin_seq_ali)
        structure_gaps = self.find_gaps(structure_seq_ali)

        ## STAGE 1 - REMOVE RESIDUES ##
        pdbin_stage1 = tempfile.NamedTemporaryFile(delete=False)
        pdbin_stage1.close()
        structure_stage1 = tempfile.NamedTemporaryFile(delete=False)
        structure_stage1.close()

        # Get first residue number to adjust list of residues to remove
        pdbin_res1     = self.residue_one(pdbin)
        structure_res1 = self.residue_one(structure)

        # Match the residue lists to fit the residue 1 number
        pdbin_gaps     = [i + structure_res1-1 for i in pdbin_gaps]
        structure_gaps = [i + pdbin_res1-1     for i in structure_gaps]

        # Use gaps of other sequence to even out
        pdb_edit.select_residues(pdbin, pdbin_stage1.name, delete=structure_gaps)
        pdb_edit.select_residues(structure, structure_stage1.name, delete=pdbin_gaps)

        ## STAGE 2 - RENUMBER RESIDUES ##
        pdb_edit.renumber_residues(pdbin_stage1.name, pdbin_mod)
        pdb_edit.renumber_residues(structure_stage1.name, structure_mod)

        os.unlink(pdbin_stage1.name)
        os.unlink(structure_stage1.name)

        return

    def residue_one(self, pdb):
        for line in open(pdb, 'r'):
            if line.startswith("ATOM"):
                line = line.split()
                return int(line[5])

    def find_gaps(self, seq):
        gaps = [i+1 for i, c in enumerate(seq) if c=="-"]
        return gaps

    def read_sequence(self, seq):
        offset = 0
        for char in seq:
            if char=="-": offset += 1
            if char!="-": break
        return offset

    def _read_list(self, list_file):
        return [l.strip() for l in open(list_file, 'r')]

    def _store(self, name, pdbin, logfile, structure, pt):
        entry = {"name": name,
                 "pdbin": pdbin,
                 "TMSCORE_log": logfile,
                 "structure": structure,
                 "tm": pt.tm,
                 "maxsub": pt.maxsub,
                 "gdtts": pt.gdtts,
                 "gdtha": pt.gdtha,
                 "rmsd": pt.rmsd,
                 "nrResiduesCommon": pt.nrResiduesCommon}
        return entry

class Statistics(object):
    def __init__(self, pickle_file, pdb_list_file):
        self.pdb_list_file = pdb_list_file
        self.pickle_file = pickle_file
        self.mean = 0.0
        self.median = 0.0
        return

    def main(self, score):
        # Load the required data
        pdb_names = self._read_list(self.pdb_list_file)
        data = pickle.load(open(self.pickle_file, 'r'))

        # Obtain a list of all defined scores for defined pdbs
        score_list = self.extract(data, pdb_names, score)
        score_list_sorted = sorted(score_list)

        assert len(score_list_sorted) > 0, \
                "No scores in score list file"

        assert len(score_list_sorted)==len(pdb_names), \
                "Divergent counts between scores and pdb names"

        self.mean   = numpy.mean(score_list_sorted)
        self.median = numpy.median(score_list_sorted)
        return

    def extract(self, data, pdb_names, score):
        return [entry[score] for entry in data \
                        if entry['name'] in pdb_names]

    def _read_list(self, list_file):
        return [os.path.basename(l.strip()).rsplit('.', 1)[0] \
                    for l in open(list_file, 'r')]

class TestTMScore(unittest.TestCase):
    def setUp(self):
        self.TM = TMscorer("foo", "bar", "cho")

    def testRead(self):
        seq1 = "-----AAAA---"
        ref1_offset = 5

        seq2 = "AAA---"
        ref2_offset = 0

        seq3 = "AAAA"
        ref3_offset = 0

        offset1 = self.TM.read_sequence(seq1)
        offset2 = self.TM.read_sequence(seq2)
        offset3 = self.TM.read_sequence(seq3)

        self.assertEqual(ref1_offset, offset1)
        self.assertEqual(ref3_offset, offset2)
        self.assertEqual(ref3_offset, offset3)

    def testGaps(self):
        seq1 = "AAAA---AA--AA"
        ref_gaps1 = [5, 6, 7, 10, 11]

        seq2 = "---AA-AA"
        ref_gaps2 = [1, 2, 3, 6]

        seq3 = "-AAA--"
        ref_gaps3 = [1, 5, 6]

        gaps1 = self.TM.find_gaps(seq1)
        gaps2 = self.TM.find_gaps(seq2)
        gaps3 = self.TM.find_gaps(seq3)

        self.assertEqual(ref_gaps1, gaps1)
        self.assertEqual(ref_gaps2, gaps2)
        self.assertEqual(ref_gaps3, gaps3)

class TestStatistics(unittest.TestCase):
    def testExtract(self):
        data = [{"name": "test1", "pdbin": "test1",
                 "TMSCORE_log": "test1", "structure": "test1",
                 "tm": 0.523, "maxsub": 0.333, "gdtts": 0.355,
                 "gdtha": 0.424, "rmsd": 1.345, "nrResiduesCommon": 10},
                {"name": "test2", "pdbin": "test2",
                 "TMSCORE_log": "test2", "structure": "test2",
                 "tm": 0.140, "maxsub": 0.222, "gdtts": 0.234,
                 "gdtha": 0.100, "rmsd": 10.444, "nrResiduesCommon": 10},
                {"name": "test3", "pdbin": "test3",
                 "TMSCORE_log": "test3", "structure": "test3",
                 "tm": 0.3, "maxsub": 0.5, "gdtts": 0.1,
                 "gdtha": 0.6, "rmsd": 1.0, "nrResiduesCommon": 10},
                {"name": "test4", "pdbin": "test4",
                 "TMSCORE_log": "test4", "structure": "test4",
                 "tm": 0.6, "maxsub": 0.8, "gdtts": 0.9,
                 "gdtha": 0.4, "rmsd": 9.345, "nrResiduesCommon": 10}]

        s = Statistics("foo", "bar")

        pdb_list = ["test1", "test3", "test4"]

        ref_data_tm = [0.523, 0.3, 0.6]
        out_data_tm = s.extract(data, pdb_list, "tm")
        self.assertItemsEqual(ref_data_tm, out_data_tm)

        ref_data_rmsd = [1.345, 1.0, 9.345]
        out_data_rmsd = s.extract(data, pdb_list, "rmsd")
        self.assertItemsEqual(ref_data_rmsd, out_data_rmsd)

        ref_data_gdtts = [0.355, 0.1, 0.9]
        out_data_gdtts = s.extract(data, pdb_list, "gdtts")
        self.assertItemsEqual(ref_data_gdtts, out_data_gdtts)

        ref_data_gdtha = [0.424, 0.6, 0.4]
        out_data_gdtha = s.extract(data, pdb_list, "gdtha")
        self.assertItemsEqual(ref_data_gdtha, out_data_gdtha)

        ref_data_maxsub = [0.333, 0.5, 0.8]
        out_data_maxsub = s.extract(data, pdb_list, "maxsub")
        self.assertItemsEqual(ref_data_maxsub, out_data_maxsub)

        pdb_list_2 = ["test2", "test4"]
        ref_data_maxsub_2 = [0.222, 0.8]
        out_data_maxsub_2 = s.extract(data, pdb_list_2, "maxsub")
        self.assertItemsEqual(ref_data_maxsub_2, out_data_maxsub_2)

if __name__ == "__main__":
    # Check whether we can import Biopythin
    if not _BIOPYTHON: sys.exit("Upgrade to a CCP4 version with the new interface to use this script")

    parser = argparse.ArgumentParser()

    parser.add_argument('-score', metavar='[ tm | maxsub | gdtts | gdtha ]',
                        type=str, default='tm', help='The score to extract [default: tm]')
    parser.add_argument("structure")
    parser.add_argument("pdb_list_file")
    parser.add_argument("tmscore")
    parser.add_argument("-wdir", type=str, default=os.getcwd())
    args = parser.parse_args()

    assert args.score in ["tm", "maxsub", "gdtts", "gdtha", "rmsd"], "Unknown score %s" % args.score

    t = TMscorer(os.path.abspath(args.structure), os.path.abspath(args.tmscore),
                os.path.abspath(args.wdir))
    t.main(args.pdb_list_file)

    pkl_file = t.pickle_file

    s = Statistics(pkl_file, os.path.abspath(args.pdb_list_file))
    s.main(args.score)
    print "Mean", s.mean
    print "Median", s.median
<commit_msg>Add root variable for when we're outside ccp4 in tmscore.py<commit_after>#!/usr/bin/env ccp4-python

'''
09.11.2015

@author: hlfsimko
'''

# System
import argparse
import collections
import logging
import os
import pickle
import sys
import tempfile
import unittest

# 3rd party
import numpy

if not "CCP4" in os.environ.keys(): raise RuntimeError('CCP4 not found')
root = os.path.join(os.environ['CCP4'], "share", "ample")
#root = os.sep.join( os.path.abspath(__file__).split( os.sep )[:-2] )
sys.path.insert(0, os.path.join(root, "parsers"))

# Custom
import ample_util
import parse_tmscore
import pdb_edit

try:
    import parse_alignment
    _BIOPYTHON = True
except ImportError:
    _BIOPYTHON = False

class TMscorer(object):
    def __init__(self, structure, tmscore_exe, wdir=None):
        self.logger = logging.getLogger()
        self.pickle_file = None
        self.structure = structure
        self.tmscore_exe = tmscore_exe
        self.workingDIR = wdir
        return

    def main(self, pdb_list_file):
        # Read all pdb files from list
        pdbs = self._read_list(pdb_list_file)

        # Compare each pdb in list against structure defined
        self.entries = self.compare_to_structure(pdbs)

        # Dumpe all data in a pickle file
        self.pickle_file = os.path.join(self.workingDIR, "tmresults.pkl")
        pickle.dump(self.entries, open(self.pickle_file, 'w'))
        return

    def compare_to_structure(self, pdb_list):
        self.logger.info('-------Evaluating decoys/models-------')

        entries = []

        for pdb in pdb_list:
            if not os.path.exists(pdb): continue

            name = os.path.basename(pdb).rsplit(".", 1)[0]
            pdbin = os.path.abspath(pdb)

            self.logger.info("Working on %s" % name)

            # Write a temporary file with the aligned residues from the native structure
            structure_mod = tempfile.NamedTemporaryFile(delete=False)
            structure_mod.close()
            pdbin_mod = tempfile.NamedTemporaryFile(delete=False)
            pdbin_mod.close()

            # Initialise the logparser that stores all the scores
            pt = parse_tmscore.TMscoreLogParser()

            # Do the try clause here to allow anything that is required from here to throw
            # exceptions. In that case we revert to the TMscoreLogParser default values of
            # 0.0 for every score.
            try:
                self.mod_structures(pdbin, pdbin_mod.name, self.structure, structure_mod.name)

                # Create a command list and execute TMscore
                log = os.path.join(self.workingDIR, name+".tmscore.log")
                cmd = [ self.tmscore_exe, pdbin_mod.name, structure_mod.name ]
                p = ample_util.run_command(cmd, logfile=log, directory=self.workingDIR)

                os.unlink(structure_mod.name)
                os.unlink(pdbin_mod.name)

                # Parse the log file to the parser
                pt.parse(log)

            except Exception:
                log = "None"

            entry = self._store(name, pdbin, log, self.structure, pt)
            entries.append(entry)

        return entries

    def mod_structures(self, pdbin, pdbin_mod, structure, structure_mod):
        ''' Make sure the decoy and the xtal pdb align to get an accurate Tm score '''

        if not _BIOPYTHON:
            raise ImportError

        # Disable the info logger to not spam the user with which chain of native extracted.
        # Happens for every model + native below
        # http://stackoverflow.com/questions/2266646/how-to-i-disable-and-re-enable-console-logging-in-python
        logging.disable( logging.INFO )

        pdbin_seq     = pdb_edit.sequence(pdbin).values()[0]
        structure_seq = pdb_edit.sequence(structure).values()[0]

        # Align the sequences to see how much of the predicted decoys are in the xtal
        aligned_seq_list = parse_alignment.AlignmentParser().align_sequences(pdbin_seq, structure_seq)
        pdbin_seq_ali     = aligned_seq_list[0]
        structure_seq_ali = aligned_seq_list[1]

        # Get the gaps in both sequences
        pdbin_gaps     = self.find_gaps(pdbin_seq_ali)
        structure_gaps = self.find_gaps(structure_seq_ali)

        ## STAGE 1 - REMOVE RESIDUES ##
        pdbin_stage1 = tempfile.NamedTemporaryFile(delete=False)
        pdbin_stage1.close()
        structure_stage1 = tempfile.NamedTemporaryFile(delete=False)
        structure_stage1.close()

        # Get first residue number to adjust list of residues to remove
        pdbin_res1     = self.residue_one(pdbin)
        structure_res1 = self.residue_one(structure)

        # Match the residue lists to fit the residue 1 number
        pdbin_gaps     = [i + structure_res1-1 for i in pdbin_gaps]
        structure_gaps = [i + pdbin_res1-1     for i in structure_gaps]

        # Use gaps of other sequence to even out
        pdb_edit.select_residues(pdbin, pdbin_stage1.name, delete=structure_gaps)
        pdb_edit.select_residues(structure, structure_stage1.name, delete=pdbin_gaps)

        ## STAGE 2 - RENUMBER RESIDUES ##
        pdb_edit.renumber_residues(pdbin_stage1.name, pdbin_mod)
        pdb_edit.renumber_residues(structure_stage1.name, structure_mod)

        os.unlink(pdbin_stage1.name)
        os.unlink(structure_stage1.name)

        return

    def residue_one(self, pdb):
        for line in open(pdb, 'r'):
            if line.startswith("ATOM"):
                line = line.split()
                return int(line[5])

    def find_gaps(self, seq):
        gaps = [i+1 for i, c in enumerate(seq) if c=="-"]
        return gaps

    def read_sequence(self, seq):
        offset = 0
        for char in seq:
            if char=="-": offset += 1
            if char!="-": break
        return offset

    def _read_list(self, list_file):
        return [l.strip() for l in open(list_file, 'r')]

    def _store(self, name, pdbin, logfile, structure, pt):
        entry = {"name": name,
                 "pdbin": pdbin,
                 "TMSCORE_log": logfile,
                 "structure": structure,
                 "tm": pt.tm,
                 "maxsub": pt.maxsub,
                 "gdtts": pt.gdtts,
                 "gdtha": pt.gdtha,
                 "rmsd": pt.rmsd,
                 "nrResiduesCommon": pt.nrResiduesCommon}
        return entry

class Statistics(object):
    def __init__(self, pickle_file, pdb_list_file):
        self.pdb_list_file = pdb_list_file
        self.pickle_file = pickle_file
        self.mean = 0.0
        self.median = 0.0
        return

    def main(self, score):
        # Load the required data
        pdb_names = self._read_list(self.pdb_list_file)
        data = pickle.load(open(self.pickle_file, 'r'))

        # Obtain a list of all defined scores for defined pdbs
        score_list = self.extract(data, pdb_names, score)
        score_list_sorted = sorted(score_list)

        assert len(score_list_sorted) > 0, \
                "No scores in score list file"

        assert len(score_list_sorted)==len(pdb_names), \
                "Divergent counts between scores and pdb names"

        self.mean   = numpy.mean(score_list_sorted)
        self.median = numpy.median(score_list_sorted)
        return

    def extract(self, data, pdb_names, score):
        return [entry[score] for entry in data \
                        if entry['name'] in pdb_names]

    def _read_list(self, list_file):
        return [os.path.basename(l.strip()).rsplit('.', 1)[0] \
                    for l in open(list_file, 'r')]

class TestTMScore(unittest.TestCase):
    def setUp(self):
        self.TM = TMscorer("foo", "bar", "cho")

    def testRead(self):
        seq1 = "-----AAAA---"
        ref1_offset = 5

        seq2 = "AAA---"
        ref2_offset = 0

        seq3 = "AAAA"
        ref3_offset = 0

        offset1 = self.TM.read_sequence(seq1)
        offset2 = self.TM.read_sequence(seq2)
        offset3 = self.TM.read_sequence(seq3)

        self.assertEqual(ref1_offset, offset1)
        self.assertEqual(ref3_offset, offset2)
        self.assertEqual(ref3_offset, offset3)

    def testGaps(self):
        seq1 = "AAAA---AA--AA"
        ref_gaps1 = [5, 6, 7, 10, 11]

        seq2 = "---AA-AA"
        ref_gaps2 = [1, 2, 3, 6]

        seq3 = "-AAA--"
        ref_gaps3 = [1, 5, 6]

        gaps1 = self.TM.find_gaps(seq1)
        gaps2 = self.TM.find_gaps(seq2)
        gaps3 = self.TM.find_gaps(seq3)

        self.assertEqual(ref_gaps1, gaps1)
        self.assertEqual(ref_gaps2, gaps2)
        self.assertEqual(ref_gaps3, gaps3)

class TestStatistics(unittest.TestCase):
    def testExtract(self):
        data = [{"name": "test1", "pdbin": "test1",
                 "TMSCORE_log": "test1", "structure": "test1",
                 "tm": 0.523, "maxsub": 0.333, "gdtts": 0.355,
                 "gdtha": 0.424, "rmsd": 1.345, "nrResiduesCommon": 10},
                {"name": "test2", "pdbin": "test2",
                 "TMSCORE_log": "test2", "structure": "test2",
                 "tm": 0.140, "maxsub": 0.222, "gdtts": 0.234,
                 "gdtha": 0.100, "rmsd": 10.444, "nrResiduesCommon": 10},
                {"name": "test3", "pdbin": "test3",
                 "TMSCORE_log": "test3", "structure": "test3",
                 "tm": 0.3, "maxsub": 0.5, "gdtts": 0.1,
                 "gdtha": 0.6, "rmsd": 1.0, "nrResiduesCommon": 10},
                {"name": "test4", "pdbin": "test4",
                 "TMSCORE_log": "test4", "structure": "test4",
                 "tm": 0.6, "maxsub": 0.8, "gdtts": 0.9,
                 "gdtha": 0.4, "rmsd": 9.345, "nrResiduesCommon": 10}]

        s = Statistics("foo", "bar")

        pdb_list = ["test1", "test3", "test4"]

        ref_data_tm = [0.523, 0.3, 0.6]
        out_data_tm = s.extract(data, pdb_list, "tm")
        self.assertItemsEqual(ref_data_tm, out_data_tm)

        ref_data_rmsd = [1.345, 1.0, 9.345]
        out_data_rmsd = s.extract(data, pdb_list, "rmsd")
        self.assertItemsEqual(ref_data_rmsd, out_data_rmsd)

        ref_data_gdtts = [0.355, 0.1, 0.9]
        out_data_gdtts = s.extract(data, pdb_list, "gdtts")
        self.assertItemsEqual(ref_data_gdtts, out_data_gdtts)

        ref_data_gdtha = [0.424, 0.6, 0.4]
        out_data_gdtha = s.extract(data, pdb_list, "gdtha")
        self.assertItemsEqual(ref_data_gdtha, out_data_gdtha)

        ref_data_maxsub = [0.333, 0.5, 0.8]
        out_data_maxsub = s.extract(data, pdb_list, "maxsub")
        self.assertItemsEqual(ref_data_maxsub, out_data_maxsub)

        pdb_list_2 = ["test2", "test4"]
        ref_data_maxsub_2 = [0.222, 0.8]
        out_data_maxsub_2 = s.extract(data, pdb_list_2, "maxsub")
        self.assertItemsEqual(ref_data_maxsub_2, out_data_maxsub_2)

if __name__ == "__main__":
    # Check whether we can import Biopythin
    if not _BIOPYTHON: sys.exit("Upgrade to a CCP4 version with the new interface to use this script")

    parser = argparse.ArgumentParser()

    parser.add_argument('-score', metavar='[ tm | maxsub | gdtts | gdtha ]',
                        type=str, default='tm', help='The score to extract [default: tm]')
    parser.add_argument("structure")
    parser.add_argument("pdb_list_file")
    parser.add_argument("tmscore")
    parser.add_argument("-wdir", type=str, default=os.getcwd())
    args = parser.parse_args()

    assert args.score in ["tm", "maxsub", "gdtts", "gdtha", "rmsd"], "Unknown score %s" % args.score

    t = TMscorer(os.path.abspath(args.structure), os.path.abspath(args.tmscore),
                os.path.abspath(args.wdir))
    t.main(args.pdb_list_file)

    pkl_file = t.pickle_file

    s = Statistics(pkl_file, os.path.abspath(args.pdb_list_file))
    s.main(args.score)
    print "Mean", s.mean
    print "Median", s.median
<|endoftext|>


<commit_before># The MIT License (MIT)
#
# Copyright (c) 2015 Eran Sandler
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from __future__ import absolute_import
import functools
import urllib2
import datetime
import getpass
import os
try:
    import simplejson as json
except ImportError:
    import json

from fabric.api import env

_web_hook_url = None

def init_slack(web_hook_url):
    global _web_hook_url
    _web_hook_url = web_hook_url

def send_slack_message(text, channel=None, username=None, web_hook_url=None):
    """
    send_slack_message - allow sending a slack message to a channel via a webhook.
    To configure a web hook:
    - Go to "Service Integrations -> Incoming WebHooks -> Add"
    - Copy the hook URL and use it via the web_hook_url function parameter or globally
      (so you won't have to repeat it) by calling init_slack once at the bottom
      of your fabfile
    - There is no longer need to worry about channel or username since the Slack
      incoming hook has already taken care of it on Slack's side, but that option
      is still available as an override
    """
    global _web_hook_url

    if not web_hook_url:
        if _web_hook_url:
            web_hook_url = _web_hook_url
        else:
            raise ValueError("web_hook_url must be set")

    data = {
        "text": text
    }

    if channel:
        # bring in the optional channel
        data["channel"] = channel

    if username:
        # bring in the optional (bot) username
        data["username"] = username

    req = urllib2.Request(web_hook_url)
    req.add_header("Content-Type", "application/json")
    urllib2.urlopen(req, json.dumps(data))

def announce_deploy(project, channel=None, username=None, web_hook_url=None):
    """
    announce_deploy - a decorator to announces a new deploy of the specificed project start,
    who started it and when. Also sends a message when the deploy ends and how long it took.
    """

    def real_announce_deploy(func):
        @functools.wraps(func)
        def inner_decorator(*args, **kwargs):
            # Get current time
            deploy_start = datetime.datetime.utcnow()
            # Get user's identity from a environment variable SLACK_IDENTITY, otherwise just get the default username
            deployment_handler = os.environ.get('SLACK_IDENTITY', getpass.getuser())
            # Send a message on deployment start...
            if env.host:
                start_message = "%s deploy started by %s on %s" % (project, deployment_handler, env.host)
            else:
                start_message = "%s deploy started by %s" % (project, deployment_handler)
            send_slack_message(start_message, channel=channel, username=username, web_hook_url=web_hook_url)
            return_value = func(*args, **kwargs)
            # ... and upon finish
            if env.host:
                end_message = "%s deploy ended by %s on %s. Took: %s" % (project, deployment_handler, env.host, str(datetime.datetime.utcnow() - deploy_start))
            else:
                end_message = "%s deploy ended by %s. Took: %s" % (project, deployment_handler, str(datetime.datetime.utcnow() - deploy_start))
            send_slack_message(end_message, channel=channel, username=username)
            return return_value
        return inner_decorator
    return real_announce_deploy

__all__ = ["init_slack", "send_slack_message", "announce_deploy"]
<commit_msg>Typo: Use a noun<commit_after># The MIT License (MIT)
#
# Copyright (c) 2015 Eran Sandler
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from __future__ import absolute_import
import functools
import urllib2
import datetime
import getpass
import os
try:
    import simplejson as json
except ImportError:
    import json

from fabric.api import env

_web_hook_url = None

def init_slack(web_hook_url):
    global _web_hook_url
    _web_hook_url = web_hook_url

def send_slack_message(text, channel=None, username=None, web_hook_url=None):
    """
    send_slack_message - allow sending a slack message to a channel via a webhook.
    To configure a web hook:
    - Go to "Service Integrations -> Incoming WebHooks -> Add"
    - Copy the hook URL and use it via the web_hook_url function parameter or globally
      (so you won't have to repeat it) by calling init_slack once at the bottom
      of your fabfile
    - There is no longer need to worry about channel or username since the Slack
      incoming hook has already taken care of it on Slack's side, but that option
      is still available as an override
    """
    global _web_hook_url

    if not web_hook_url:
        if _web_hook_url:
            web_hook_url = _web_hook_url
        else:
            raise ValueError("web_hook_url must be set")

    data = {
        "text": text
    }

    if channel:
        # bring in the optional channel
        data["channel"] = channel

    if username:
        # bring in the optional (bot) username
        data["username"] = username

    req = urllib2.Request(web_hook_url)
    req.add_header("Content-Type", "application/json")
    urllib2.urlopen(req, json.dumps(data))

def announce_deploy(project, channel=None, username=None, web_hook_url=None):
    """
    announce_deploy - a decorator to announces a new deploy of the specificed project start,
    who started it and when. Also sends a message when the deploy ends and how long it took.
    """

    def real_announce_deploy(func):
        @functools.wraps(func)
        def inner_decorator(*args, **kwargs):
            # Get current time
            deploy_start = datetime.datetime.utcnow()
            # Get user's identity from a environment variable SLACK_IDENTITY, otherwise just get the default username
            deployment_handler = os.environ.get('SLACK_IDENTITY', getpass.getuser())
            # Send a message on deployment start...
            if env.host:
                start_message = "%s deployment started by %s on %s" % (project, deployment_handler, env.host)
            else:
                start_message = "%s deployment started by %s" % (project, deployment_handler)
            send_slack_message(start_message, channel=channel, username=username, web_hook_url=web_hook_url)
            return_value = func(*args, **kwargs)
            # ... and upon finish
            if env.host:
                end_message = "%s deployment ended by %s on %s. Took: %s" % (project, deployment_handler, env.host, str(datetime.datetime.utcnow() - deploy_start))
            else:
                end_message = "%s deployment ended by %s. Took: %s" % (project, deployment_handler, str(datetime.datetime.utcnow() - deploy_start))
            send_slack_message(end_message, channel=channel, username=username)
            return return_value
        return inner_decorator
    return real_announce_deploy

__all__ = ["init_slack", "send_slack_message", "announce_deploy"]
<|endoftext|>


<commit_before>import numpy as np
import numpy.ma as ma
from numpy import linalg as LA
import matplotlib.pyplot as plt
import itertools
import collections
from scipy import stats

def acf(x, lags=500, exclude=None):
    if exclude is None:
        exclude = np.zeros(x.shape)
    exclude = np.cumsum(exclude.astype(int))

    # from stackexchange
    x = x - x.mean()  # remove mean
    if type(lags) is int:
        lags = range(lags)

    C = ma.zeros((len(lags),))
    for i, l in enumerate(lags):
        if l == 0:
            C[i] = 1
        else:
            x0 = x[:-l].copy()
            x1 = x[l:].copy()
            reject = (exclude[l:]-exclude[:-l])>0
            x0[reject] = ma.masked
            x1[reject] = ma.masked
            C[i] = ma.corrcoef(x0, x1)[0, 1]
    return C

def dotacf(x, lags=500, exclude=None):
    if exclude is None:
        exclude = np.zeros(x.shape)
    exclude = np.cumsum(exclude.astype(int))

    if type(lags) is int:
        lags = xrange(lags)
    C = ma.zeros((len(lags),))
    for i, l in enumerate(lags):
        if l == 0:
            C[i] = (x*x).sum(axis=1).mean()
        else:
            x0 = x[:-l, :].copy()
            x1 = x[l:, :].copy()
            reject = (exclude[l:]-exclude[:-l])>0
            x0[reject, :] = ma.masked
            x1[reject, :] = ma.masked
            C[i] = (x0*x1).sum(axis=1).mean()
    return C

def drift(x, lags=500, exclude=None):
    if exclude is None:
        exclude = np.zeros(x.shape)
    exclude = np.cumsum(exclude.astype(int))

    if type(lags) is int:
        lags = xrange(lags)
    mu = ma.zeros((len(lags),))
    for i, lag in enumerate(lags):
        if lag==0:
            mu[i] = 0
        else:
            x0 = x[lag:].copy()
            x1 = x[:-lag].copy()
            reject = (exclude[lag:]-exclude[:-lag])>0
            x0[reject] = ma.masked
            x1[reject] = ma.masked
            displacements = x0 - x1
            mu[i] = displacements.mean()
    return mu

def unwrapma(x):
    # Adapted from numpy unwrap, this version ignores missing data
    idx = ma.array(np.arange(0,x.shape[0]), mask=x.mask)
    idxc = idx.compressed()
    xc = x.compressed()
    dd = np.diff(xc)
    ddmod = np.mod(dd+np.pi, 2*np.pi)-np.pi
    ddmod[(ddmod==-np.pi) & (dd > 0)] = np.pi
    phc_correct = ddmod - dd
    phc_correct[np.abs(dd)<np.pi] = 0
    ph_correct = np.zeros(x.shape)
    ph_correct[idxc[1:]] = phc_correct
    up = x + ph_correct.cumsum()
    return up

def nextpow2(n):
    '''
    Returns the next highest power of 2 from n
    '''
    m_f = np.log2(n)
    m_i = np.ceil(m_f)
    return 2**m_i

def phaserand(X, independent=False):
    '''
    Generates a randomized surrogate dataset for X, preserving linear temporal
    correlations. If independent is False (default), linear correlations
    between columns of x are also preserved.

    If X contains missing values, they are filled with the mean of that
    channel.

    The algorithm works by randomizing the phases in the Fourier domain. For
    non-independent shuffling, the same random phases are used for each
    channel.

    References:
    Theiler, J., Eubank, S., Longtin, A., Galdrikian, B., & Doyne Farmer, J.
        (1992). Testing for nonlinearity in time series: the method of
        surrogate data. Physica D: Nonlinear Phenomena, 58(1), 77-94.
    Prichard, D. and Theiler, J. (1994) Generating surrogate data for time
        series with several simultaneously measured variables. Phys. Rev.
        Lett. 73(7), 951-954.
    Podobnik, B., Fu, D. F., Stanley, H. E., & Ivanov, P. C. (2007).
        Power-law autocorrelated stochastic processes with long-range
        cross-correlations. The European Physical Journal B, 56(1), 47-52.
    '''
    # Deal with array vs matrix by adding new axis
    if len(X.shape) == 1:
        X = X[:, np.newaxis]

    # Deal with missing data
    if isinstance(X, ma.MaskedArray):
        X = X.filled(X.mean(axis=0))  # fill missing values with the mean

    # Fourier transform and extract amplitude and phases
    # The frequencies are shifted so 0 is centered (fftshift)
    N = int(nextpow2(X.shape[0]))  # size for FFT
    h = np.floor(N/2)  # half the length of the data
    Z = np.fft.fft(X, N, axis=0)
    M = np.fft.fftshift(np.abs(Z), axes=0)  # the amplitudes
    phase = np.fft.fftshift(np.angle(Z), axes=0)  # the original phases

    # Randomize the phases. The phases need to be symmetric for postivie and
    # negative frequencies.
    if independent:  # generate random phases for each channel
        randphase = 2.*np.pi*np.random.rand((h-1, X.shape[1]))  # random phases
        newphase = np.zeros((N, X.shape[1]))  # new phases to use
        newphase[0, :] = phase[0, :]  # keep the zero freq (don't know why)
        newphase[1:h, :] = randphase[::-1, :]
        newphase[h, :] = phase[h, :]
        newphase[h+1:, :] = -randphase
    else:  # generate one set of random phases (same as above)
        randphase = 2.*np.pi*np.random.rand(h-1)
        newphase = np.zeros((N, X.shape[1]))
        newphase[0, :] = phase[0, :]
        newphase[1:h, :] = randphase[::-1, np.newaxis]
        newphase[h, :] = phase[h, :]
        newphase[h+1:, :] = -randphase[:, np.newaxis]

    # Reconstruct the signal from the original amplitude and the new phases
    z2 = M*np.exp(newphase*1.j)

    # Return the time-domain signal (have to undo the fftshift and keep the
    # orignal length)
    return np.fft.ifft(np.fft.ifftshift(z2, axes=0),
                       axis=0).real[:X.shape[0], :].squeeze()
<commit_msg>Old updates for AR modeling/phase randomization<commit_after>import numpy as np
import numpy.ma as ma
from numpy import linalg as LA
import matplotlib.pyplot as plt
import itertools
import collections
from scipy import stats

def acf(x, lags=500, exclude=None):
    if exclude is None:
        exclude = np.zeros(x.shape)
    exclude = np.cumsum(exclude.astype(int))

    # from stackexchange
    x = x - x.mean()  # remove mean
    if type(lags) is int:
        lags = range(lags)

    C = ma.zeros((len(lags),))
    for i, l in enumerate(lags):
        if l == 0:
            C[i] = 1
        else:
            x0 = x[:-l].copy()
            x1 = x[l:].copy()
            reject = (exclude[l:]-exclude[:-l])>0
            x0[reject] = ma.masked
            x1[reject] = ma.masked
            C[i] = ma.corrcoef(x0, x1)[0, 1]
    return C

def dotacf(x, lags=500, exclude=None):
    if exclude is None:
        exclude = np.zeros(x.shape)
    exclude = np.cumsum(exclude.astype(int))

    if type(lags) is int:
        lags = xrange(lags)
    C = ma.zeros((len(lags),))
    for i, l in enumerate(lags):
        if l == 0:
            C[i] = (x*x).sum(axis=1).mean()
        else:
            x0 = x[:-l, :].copy()
            x1 = x[l:, :].copy()
            reject = (exclude[l:]-exclude[:-l])>0
            x0[reject, :] = ma.masked
            x1[reject, :] = ma.masked
            C[i] = (x0*x1).sum(axis=1).mean()
    return C

def pacfe(p,j,List):
    '''
    Partial autocorrelation function estimates
    p is the order of the AR(p) process
    j is the coefficient in an AR(p) process
    '''
    if p==2 and j==1:
        return (acf(j,List)*(1-acf(p,List)))/(1-(acf(j,List))**2)
    elif p==2 and j==2:
        return (acf(2,List)-(acf(1,List))**2)/(1-(acf(1,List))**2)
    elif p==j and p!=2 and j!=2:
        c=0
        for a in range(1,p):
            c+=pacfe(p-1,a,List)*acf(p-a,List)
        d=0
        for b in range(1,p):
            d+=pacfe(p-1,b,List)*acf(b,List)
        return (acf(p,List)-c)/(1-d)
    else:
        return pacfe(p-1,j,List)-pacfe(p,p,List)*pacfe(p-1,p-j,List)

def drift(x, lags=500, exclude=None):
    if exclude is None:
        exclude = np.zeros(x.shape)
    exclude = np.cumsum(exclude.astype(int))

    if type(lags) is int:
        lags = xrange(lags)
    mu = ma.zeros((len(lags),))
    for i, lag in enumerate(lags):
        if lag==0:
            mu[i] = 0
        else:
            x0 = x[lag:].copy()
            x1 = x[:-lag].copy()
            reject = (exclude[lag:]-exclude[:-lag])>0
            x0[reject] = ma.masked
            x1[reject] = ma.masked
            displacements = x0 - x1
            mu[i] = displacements.mean()
    return mu

def unwrapma(x):
    # Adapted from numpy unwrap, this version ignores missing data
    idx = ma.array(np.arange(0,x.shape[0]), mask=x.mask)
    idxc = idx.compressed()
    xc = x.compressed()
    dd = np.diff(xc)
    ddmod = np.mod(dd+np.pi, 2*np.pi)-np.pi
    ddmod[(ddmod==-np.pi) & (dd > 0)] = np.pi
    phc_correct = ddmod - dd
    phc_correct[np.abs(dd)<np.pi] = 0
    ph_correct = np.zeros(x.shape)
    ph_correct[idxc[1:]] = phc_correct
    up = x + ph_correct.cumsum()
    return up

def nextpow2(n):
    '''
    Returns the next highest power of 2 from n
    '''
    m_f = np.log2(n)
    m_i = np.ceil(m_f)
    return 2**m_i

def phaserand(X, independent=False, reduceHighFreqNoise=True):
    '''
    Generates a randomized surrogate dataset for X, preserving linear temporal
    correlations. If independent is False (default), linear correlations
    between columns of x are also preserved.

    If X contains missing values, they are filled with the mean of that
    channel.

    The algorithm works by randomizing the phases in the Fourier domain. For
    non-independent shuffling, the same random phases are used for each
    channel.

    References:
    Theiler, J., Eubank, S., Longtin, A., Galdrikian, B., & Doyne Farmer, J.
        (1992). Testing for nonlinearity in time series: the method of
        surrogate data. Physica D: Nonlinear Phenomena, 58(1), 77-94.
    Prichard, D. and Theiler, J. (1994) Generating surrogate data for time
        series with several simultaneously measured variables. Phys. Rev.
        Lett. 73(7), 951-954.
    Podobnik, B., Fu, D. F., Stanley, H. E., & Ivanov, P. C. (2007).
        Power-law autocorrelated stochastic processes with long-range
        cross-correlations. The European Physical Journal B, 56(1), 47-52.
    '''
    # Deal with array vs matrix by adding new axis
    if len(X.shape) == 1:
        X = X[:, np.newaxis]

    # Deal with missing data
    if isinstance(X, ma.MaskedArray):
        # truncate all missing data at beginning and end
        idxNotAllMissing = (~np.all(X.mask, axis=1)).nonzero()[0]
        X = X[idxNotAllMissing[0]:idxNotAllMissing[-1], :]
        X = X.filled(X.mean(axis=0))  # fill interior mask with the mean

    # Reduce high-frequency noise by min difference between first and last
    if reduceHighFreqNoise:
        delta = X - X[0, :]
        threshold = 1e-3*np.std(X, axis=0)
        # find last pt in which all the channels are about the same as the beginning
        # and also the index is even
        goodEndPt = np.nonzero((np.all(np.abs(delta) < threshold, axis=1)) &
                               (np.arange(0, X.shape[1]) % 2 == 0))[0][-1]
        if goodEndPt > X.shape[0]/2:  # make sure we keep at least half the data
            X = X[:goodEndPt, :]

    # Fourier transform and extract amplitude and phases
    # The frequencies are shifted so 0 is centered (fftshift)
    N = X.shape[0] #int(nextpow2(X.shape[0]))  # size for FFT
    if N % 2 != 0:
        N = N-1
    h = np.floor(N/2)  # half the length of the data
    Z = np.fft.fft(X, N, axis=0)
    M = np.fft.fftshift(np.abs(Z), axes=0)  # the amplitudes
    phase = np.fft.fftshift(np.angle(Z), axes=0)  # the original phases

    # Randomize the phases. The phases need to be symmetric for postivie and
    # negative frequencies.
    if independent:  # generate random phases for each channel
        randphase = 2.*np.pi*np.random.rand((h-1, X.shape[1]))  # random phases
        newphase = np.zeros((N, X.shape[1]))  # new phases to use
        newphase[0, :] = phase[0, :]  # keep the zero freq (don't know why)
        newphase[1:h, :] = randphase[::-1, :]
        newphase[h, :] = phase[h, :]
        newphase[h+1:, :] = -randphase
    else:  # generate one set of random phases (same as above)
        randphase = 2.*np.pi*np.random.rand(h-1)
        newphase = np.zeros((N, X.shape[1]))
        newphase[0, :] = phase[0, :]
        newphase[1:h, :] = randphase[::-1, np.newaxis]
        newphase[h, :] = phase[h, :]
        newphase[h+1:, :] = -randphase[:, np.newaxis]

    # Reconstruct the signal from the original amplitude and the new phases
    z2 = M*np.exp(newphase*1.j)

    # Return the time-domain signal
    return np.fft.ifft(np.fft.ifftshift(z2, axes=0),
                       axis=0).real.squeeze()
<|endoftext|>