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

"""Functions to load data from Worldbank API.

We use its bulkdownload utilities.

Source link: `WorldBank website`_

.. _WorldBank website: http://data.worldbank.org

"""

import os.path as osp

import logging
import requests
import pandas as pd

from urllib.parse import urlsplit
from concurrent.futures import ThreadPoolExecutor, as_completed

from . common import DataFactory, download

logger = logging.getLogger("WorldBankLoader")

class WorldBankLoader(DataFactory):
    __doc__ = """T.B.D"""
    url = 'http://api.worldbank.org/v2/datacatalog?format=json'

    def load_metadata(self):
        def load_url(url, page=None):
            if page:
                url = url + '&page=' + str(page)
            res = requests.get(url)
            return res.json()

        def handle_data(data):
            for v in data['datacatalog']:
                d = {}
                for v_ in v['metatype']:
                    d[v_['id']] = v_['value']
                yield d

        first_page = self.url
        data = load_url(first_page)

        metadata = []

        for d in handle_data(data):
            metadata.append(d)

        pages = range(2, data['pages']+1)

        executor = ThreadPoolExecutor(max_workers=3)
        res = [executor.submit(load_url, self.url, page=p) for p in pages]

        for r in as_completed(res):
            for d in handle_data(r.result()):
                # print(d)
                metadata.append(d)

        metadata = pd.DataFrame.from_records(metadata)
        self.metadata = metadata
        return metadata

    def has_newer_source(self, dataset, date):
        if self.metadata is None:
            self.load_metadata()
        metadata = self.metadata
        lastrevisiondate = metadata.loc[
            metadata.acronym == dataset,
            'lastrevisiondate'].values[0]
        if pd.to_datetime(lastrevisiondate) > pd.to_datetime(date):
            return True
        else:
            return False

    def bulk_download(self, dataset, out_dir):
        if self.metadata is None:
            self.load_metadata()
        metadata = self.metadata
        s = metadata.loc[metadata.acronym == dataset, 'bulkdownload'].values[0]
        url = s.split(';')[1].split('=')[1]
        fs_path = osp.join(out_dir, osp.basename(urlsplit(url).path))
        download(url, fs_path)

        return out_dir
<commit_msg>WB downloader: pass parameters to download function<commit_after># -*- coding: utf-8 -*-

"""Functions to load data from Worldbank API.

We use its bulkdownload utilities.

Source link: `WorldBank website`_

.. _WorldBank website: http://data.worldbank.org

"""

import os.path as osp

import logging
import requests
import pandas as pd

from urllib.parse import urlsplit
from concurrent.futures import ThreadPoolExecutor, as_completed

from . common import DataFactory, download

logger = logging.getLogger("WorldBankLoader")

class WorldBankLoader(DataFactory):
    __doc__ = """T.B.D"""
    url = 'http://api.worldbank.org/v2/datacatalog?format=json'

    def load_metadata(self):
        def load_url(url, page=None):
            if page:
                url = url + '&page=' + str(page)
            res = requests.get(url)
            return res.json()

        def handle_data(data):
            for v in data['datacatalog']:
                d = {}
                for v_ in v['metatype']:
                    d[v_['id']] = v_['value']
                yield d

        first_page = self.url
        data = load_url(first_page)

        metadata = []

        for d in handle_data(data):
            metadata.append(d)

        pages = range(2, data['pages']+1)

        executor = ThreadPoolExecutor(max_workers=3)
        res = [executor.submit(load_url, self.url, page=p) for p in pages]

        for r in as_completed(res):
            for d in handle_data(r.result()):
                # print(d)
                metadata.append(d)

        metadata = pd.DataFrame.from_records(metadata)
        self.metadata = metadata
        return metadata

    def has_newer_source(self, dataset, date):
        if self.metadata is None:
            self.load_metadata()
        metadata = self.metadata
        lastrevisiondate = metadata.loc[
            metadata.acronym == dataset,
            'lastrevisiondate'].values[0]
        if pd.to_datetime(lastrevisiondate) > pd.to_datetime(date):
            return True
        else:
            return False

    def bulk_download(self, dataset, out_dir, **kwargs):
        if self.metadata is None:
            self.load_metadata()
        metadata = self.metadata
        s = metadata.loc[metadata.acronym == dataset, 'bulkdownload'].values[0]
        url = s.split(';')[1].split('=')[1]
        fs_path = osp.join(out_dir, osp.basename(urlsplit(url).path))
        download(url, fs_path, **kwargs)

        return out_dir
<|endoftext|>


<commit_before>"""
Progress bar implementation on top of prompt_toolkit.

::

    with progress_bar(...) as pb:
        for item in pb(data):
            ...
"""
from __future__ import unicode_literals
from prompt_toolkit.application import Application
from prompt_toolkit.eventloop import get_event_loop
from prompt_toolkit.filters import Condition, is_done, renderer_height_is_known
from prompt_toolkit.formatted_text import to_formatted_text
from prompt_toolkit.input.defaults import get_default_input
from prompt_toolkit.key_binding import KeyBindings
from prompt_toolkit.layout import Layout, Window, ConditionalContainer, FormattedTextControl, HSplit, VSplit
from prompt_toolkit.layout.controls import UIControl, UIContent
from prompt_toolkit.layout.dimension import D
from prompt_toolkit.output.defaults import create_output
from prompt_toolkit.styles import BaseStyle
from prompt_toolkit.utils import in_main_thread

import functools
import contextlib
import datetime
import os
import signal
import six
import threading
import time
import traceback
import sys

from . import formatters as f

__all__ = (
    'progress_bar',
)

def create_key_bindings():
    """
    Key bindings handled by the progress bar.
    (The main thread is not supposed to handle any key bindings.)
    """
    kb = KeyBindings()

    @kb.add('c-l')
    def _(event):
        event.app.renderer.clear()

    @kb.add('c-c')
    def _(event):
        # Send KeyboardInterrupt to the main thread.
        os.kill(os.getpid(), signal.SIGINT)

    return kb

def create_default_formatters():
    return [
        f.Label(),
        f.Text(' '),
        f.Percentage(),
        f.Text(' '),
        f.Bar(),
        f.Text(' '),
        f.Progress(),
        f.Text(' '),
        f.Text('eta [', style='class:time-left'),
        f.TimeLeft(),
        f.Text(']', style='class:time-left'),
        f.Text(' '),
    ]

class progress_bar(object):
    """
    Progress bar context manager.

    Usage ::

        with progress_bar(...) as pb:
            for item in pb(data):
                ...

    :param title: Text to be displayed above the progress bars. This can be a
        callable or formatted text as well.
    :param formatters: List of `Formatter` instances.
    :param bottom_toolbar: Text to be displayed in the bottom toolbar.
        This can be a callable or formatted text.
    :param style: `prompt_toolkit` ``Style`` instance.
    :param key_bindings: `KeyBindings` instance.
    :param file: The file object used for rendering, by default `sys.stderr` is used.

    :param output: `prompt_toolkit` `Output` instance.
    :param input: `prompt_toolkit` `Input` instance.
    """
    def __init__(self, title=None, formatters=None, bottom_toolbar=None,
                 style=None, key_bindings=None, file=None, output=None, input=None):
        assert formatters is None or (
            isinstance(formatters, list) and all(isinstance(fo, f.Formatter) for fo in formatters))
        assert style is None or isinstance(style, BaseStyle)
        assert key_bindings is None or isinstance(key_bindings, KeyBindings)

        self.title = title
        self.formatters = formatters or create_default_formatters()
        self.bottom_toolbar = bottom_toolbar
        self.counters = []
        self.style = style
        self.key_bindings = key_bindings

        self.output = output or create_output(stdout=file or sys.stderr)
        self.input = input or get_default_input()

        self._thread = None

        self._loop = get_event_loop()
        self._previous_winch_handler = None
        self._has_sigwinch = False

    def __enter__(self):
        # Create UI Application.
        title_toolbar = ConditionalContainer(
            Window(FormattedTextControl(lambda: self.title), height=1, style='class:progressbar,title'),
            filter=Condition(lambda: self.title is not None))

        bottom_toolbar = ConditionalContainer(
            Window(FormattedTextControl(lambda: self.bottom_toolbar,
                                        style='class:bottom-toolbar.text'),
                   style='class:bottom-toolbar',
                   height=1),
            filter=~is_done & renderer_height_is_known &
                Condition(lambda: self.bottom_toolbar is not None))

        def width_for_formatter(formatter):
            # Needs to be passed as callable (partial) to the 'width'
            # parameter, because we want to call it on every resize.
            return formatter.get_width(progress_bar=self)

        progress_controls = [
            Window(content=_ProgressControl(self, f), width=functools.partial(width_for_formatter, f))
            for f in self.formatters
        ]

        self.app = Application(
            min_redraw_interval=.05,
            layout=Layout(HSplit([
                title_toolbar,
                VSplit(progress_controls,
                       height=lambda: D(
                           preferred=len(self.counters),
                           max=len(self.counters))),
                Window(),
                bottom_toolbar,
            ])),
            style=self.style,
            key_bindings=self.key_bindings,
            output=self.output,
            input=self.input)

        # Run application in different thread.
        def run():
            with _auto_refresh_context(self.app, .3):
                try:
                    self.app.run()
                except Exception as e:
                    traceback.print_exc()
                    print(e)

        self._thread = threading.Thread(target=run)
        self._thread.start()

        # Attach WINCH signal handler in main thread.
        # (Interrupt that we receive during resize events.)
        self._has_sigwinch = hasattr(signal, 'SIGWINCH') and in_main_thread()
        if self._has_sigwinch:
            self._previous_winch_handler = self._loop.add_signal_handler(
                signal.SIGWINCH, self.app.invalidate)

        return self

    def __exit__(self, *a):
        # Quit UI application.
        if self.app.is_running:
            self.app.set_result(None)

        # Remove WINCH handler.
        if self._has_sigwinch:
            self._loop.add_signal_handler(signal.SIGWINCH, self._previous_winch_handler)

        self._thread.join()

    def __call__(self, data=None, label='', remove_when_done=False, total=None):
        """
        Start a new counter.

        :param label: Title text or description for this progress.
        :param remove_when_done: When `True`, hide this progress bar.
        :param total: Specify the maximum value if it can't be calculated by
            calling ``len``.
        """
        assert isinstance(label, six.text_type)
        assert isinstance(remove_when_done, bool)

        counter = ProgressBarCounter(
            self, data, label=label, remove_when_done=remove_when_done, total=total)
        self.counters.append(counter)
        return counter

    def invalidate(self):
        self.app.invalidate()

class _ProgressControl(UIControl):
    """
    User control for the progress bar.
    """
    def __init__(self, progress_bar, formatter):
        self.progress_bar = progress_bar
        self.formatter = formatter
        self._key_bindings = create_key_bindings()

    def create_content(self, width, height):
        items = []
        for pr in self.progress_bar.counters:
            try:
                text = self.formatter.format(self.progress_bar, pr, width)
            except BaseException:
                traceback.print_exc()
                text = 'ERROR'

            items.append(to_formatted_text(text))

        def get_line(i):
            return items[i]

        return UIContent(
            get_line=get_line,
            line_count=len(items),
            show_cursor=False)

    def is_focussable(self):
        return True  # Make sure that the key bindings work.

    def get_key_bindings(self):
        return self._key_bindings

class ProgressBarCounter(object):
    """
    An individual counter (A progress bar can have multiple counters).
    """
    def __init__(self, progress_bar, data=None, label='', remove_when_done=False, total=None):
        self.start_time = datetime.datetime.now()
        self.progress_bar = progress_bar
        self.data = data
        self.current = 0
        self.label = label
        self.remove_when_done = remove_when_done
        self.done = False

        if total is None:
            try:
                self.total = len(data)
            except TypeError:
                self.total = None  # We don't know the total length.
        else:
            self.total = total

    def __iter__(self):
        try:
            for item in self.data:
                self.current += 1
                self.progress_bar.invalidate()
                yield item
        finally:
            self.done = True

            if self.remove_when_done:
                self.progress_bar.counters.remove(self)

    @property
    def percentage(self):
        if self.total is None:
            return 0
        else:
            return self.current * 100 / max(self.total, 1)

    @property
    def time_elapsed(self):
        """
        return how much time has been elapsed since the start.
        """
        return datetime.datetime.now() - self.start_time

    @property
    def time_left(self):
        """
        Timedelta representing the time left.
        """
        if self.total is None:
            return None
        else:
            return self.time_elapsed * (100 - self.percentage) / self.percentage

@contextlib.contextmanager
def _auto_refresh_context(app, refresh_interval=None):
    " Return a context manager for the auto-refresh loop. "
    done = [False]  # nonlocal

    # Enter.

    def run():
        while not done[0]:
            time.sleep(refresh_interval)
            app.invalidate()

    if refresh_interval:
        t = threading.Thread(target=run)
        t.daemon = True
        t.start()

    try:
        yield
    finally:
        # Exit.
        done[0] = True
<commit_msg>Use __stderr__ by default in shortcuts.progress_bar. Fixes bug when using patch_stdout.<commit_after>"""
Progress bar implementation on top of prompt_toolkit.

::

    with progress_bar(...) as pb:
        for item in pb(data):
            ...
"""
from __future__ import unicode_literals
from prompt_toolkit.application import Application
from prompt_toolkit.eventloop import get_event_loop
from prompt_toolkit.filters import Condition, is_done, renderer_height_is_known
from prompt_toolkit.formatted_text import to_formatted_text
from prompt_toolkit.input.defaults import get_default_input
from prompt_toolkit.key_binding import KeyBindings
from prompt_toolkit.layout import Layout, Window, ConditionalContainer, FormattedTextControl, HSplit, VSplit
from prompt_toolkit.layout.controls import UIControl, UIContent
from prompt_toolkit.layout.dimension import D
from prompt_toolkit.output.defaults import create_output
from prompt_toolkit.styles import BaseStyle
from prompt_toolkit.utils import in_main_thread

import functools
import contextlib
import datetime
import os
import signal
import six
import threading
import time
import traceback
import sys

from . import formatters as f

__all__ = (
    'progress_bar',
)

def create_key_bindings():
    """
    Key bindings handled by the progress bar.
    (The main thread is not supposed to handle any key bindings.)
    """
    kb = KeyBindings()

    @kb.add('c-l')
    def _(event):
        event.app.renderer.clear()

    @kb.add('c-c')
    def _(event):
        # Send KeyboardInterrupt to the main thread.
        os.kill(os.getpid(), signal.SIGINT)

    return kb

def create_default_formatters():
    return [
        f.Label(),
        f.Text(' '),
        f.Percentage(),
        f.Text(' '),
        f.Bar(),
        f.Text(' '),
        f.Progress(),
        f.Text(' '),
        f.Text('eta [', style='class:time-left'),
        f.TimeLeft(),
        f.Text(']', style='class:time-left'),
        f.Text(' '),
    ]

class progress_bar(object):
    """
    Progress bar context manager.

    Usage ::

        with progress_bar(...) as pb:
            for item in pb(data):
                ...

    :param title: Text to be displayed above the progress bars. This can be a
        callable or formatted text as well.
    :param formatters: List of `Formatter` instances.
    :param bottom_toolbar: Text to be displayed in the bottom toolbar.
        This can be a callable or formatted text.
    :param style: `prompt_toolkit` ``Style`` instance.
    :param key_bindings: `KeyBindings` instance.
    :param file: The file object used for rendering, by default `sys.stderr` is used.

    :param output: `prompt_toolkit` `Output` instance.
    :param input: `prompt_toolkit` `Input` instance.
    """
    def __init__(self, title=None, formatters=None, bottom_toolbar=None,
                 style=None, key_bindings=None, file=None, output=None, input=None):
        assert formatters is None or (
            isinstance(formatters, list) and all(isinstance(fo, f.Formatter) for fo in formatters))
        assert style is None or isinstance(style, BaseStyle)
        assert key_bindings is None or isinstance(key_bindings, KeyBindings)

        self.title = title
        self.formatters = formatters or create_default_formatters()
        self.bottom_toolbar = bottom_toolbar
        self.counters = []
        self.style = style
        self.key_bindings = key_bindings

        # Note that we use __stderr__ as default error output, because that
        # works best with `patch_stdout`.
        self.output = output or create_output(stdout=file or sys.__stderr__)
        self.input = input or get_default_input()

        self._thread = None

        self._loop = get_event_loop()
        self._previous_winch_handler = None
        self._has_sigwinch = False

    def __enter__(self):
        # Create UI Application.
        title_toolbar = ConditionalContainer(
            Window(FormattedTextControl(lambda: self.title), height=1, style='class:progressbar,title'),
            filter=Condition(lambda: self.title is not None))

        bottom_toolbar = ConditionalContainer(
            Window(FormattedTextControl(lambda: self.bottom_toolbar,
                                        style='class:bottom-toolbar.text'),
                   style='class:bottom-toolbar',
                   height=1),
            filter=~is_done & renderer_height_is_known &
                Condition(lambda: self.bottom_toolbar is not None))

        def width_for_formatter(formatter):
            # Needs to be passed as callable (partial) to the 'width'
            # parameter, because we want to call it on every resize.
            return formatter.get_width(progress_bar=self)

        progress_controls = [
            Window(content=_ProgressControl(self, f), width=functools.partial(width_for_formatter, f))
            for f in self.formatters
        ]

        self.app = Application(
            min_redraw_interval=.05,
            layout=Layout(HSplit([
                title_toolbar,
                VSplit(progress_controls,
                       height=lambda: D(
                           preferred=len(self.counters),
                           max=len(self.counters))),
                Window(),
                bottom_toolbar,
            ])),
            style=self.style,
            key_bindings=self.key_bindings,
            output=self.output,
            input=self.input)

        # Run application in different thread.
        def run():
            with _auto_refresh_context(self.app, .3):
                try:
                    self.app.run()
                except Exception as e:
                    traceback.print_exc()
                    print(e)

        self._thread = threading.Thread(target=run)
        self._thread.start()

        # Attach WINCH signal handler in main thread.
        # (Interrupt that we receive during resize events.)
        self._has_sigwinch = hasattr(signal, 'SIGWINCH') and in_main_thread()
        if self._has_sigwinch:
            self._previous_winch_handler = self._loop.add_signal_handler(
                signal.SIGWINCH, self.app.invalidate)

        return self

    def __exit__(self, *a):
        # Quit UI application.
        if self.app.is_running:
            self.app.set_result(None)

        # Remove WINCH handler.
        if self._has_sigwinch:
            self._loop.add_signal_handler(signal.SIGWINCH, self._previous_winch_handler)

        self._thread.join()

    def __call__(self, data=None, label='', remove_when_done=False, total=None):
        """
        Start a new counter.

        :param label: Title text or description for this progress.
        :param remove_when_done: When `True`, hide this progress bar.
        :param total: Specify the maximum value if it can't be calculated by
            calling ``len``.
        """
        assert isinstance(label, six.text_type)
        assert isinstance(remove_when_done, bool)

        counter = ProgressBarCounter(
            self, data, label=label, remove_when_done=remove_when_done, total=total)
        self.counters.append(counter)
        return counter

    def invalidate(self):
        self.app.invalidate()

class _ProgressControl(UIControl):
    """
    User control for the progress bar.
    """
    def __init__(self, progress_bar, formatter):
        self.progress_bar = progress_bar
        self.formatter = formatter
        self._key_bindings = create_key_bindings()

    def create_content(self, width, height):
        items = []
        for pr in self.progress_bar.counters:
            try:
                text = self.formatter.format(self.progress_bar, pr, width)
            except BaseException:
                traceback.print_exc()
                text = 'ERROR'

            items.append(to_formatted_text(text))

        def get_line(i):
            return items[i]

        return UIContent(
            get_line=get_line,
            line_count=len(items),
            show_cursor=False)

    def is_focussable(self):
        return True  # Make sure that the key bindings work.

    def get_key_bindings(self):
        return self._key_bindings

class ProgressBarCounter(object):
    """
    An individual counter (A progress bar can have multiple counters).
    """
    def __init__(self, progress_bar, data=None, label='', remove_when_done=False, total=None):
        self.start_time = datetime.datetime.now()
        self.progress_bar = progress_bar
        self.data = data
        self.current = 0
        self.label = label
        self.remove_when_done = remove_when_done
        self.done = False

        if total is None:
            try:
                self.total = len(data)
            except TypeError:
                self.total = None  # We don't know the total length.
        else:
            self.total = total

    def __iter__(self):
        try:
            for item in self.data:
                self.current += 1
                self.progress_bar.invalidate()
                yield item
        finally:
            self.done = True

            if self.remove_when_done:
                self.progress_bar.counters.remove(self)

    @property
    def percentage(self):
        if self.total is None:
            return 0
        else:
            return self.current * 100 / max(self.total, 1)

    @property
    def time_elapsed(self):
        """
        return how much time has been elapsed since the start.
        """
        return datetime.datetime.now() - self.start_time

    @property
    def time_left(self):
        """
        Timedelta representing the time left.
        """
        if self.total is None:
            return None
        else:
            return self.time_elapsed * (100 - self.percentage) / self.percentage

@contextlib.contextmanager
def _auto_refresh_context(app, refresh_interval=None):
    " Return a context manager for the auto-refresh loop. "
    done = [False]  # nonlocal

    # Enter.

    def run():
        while not done[0]:
            time.sleep(refresh_interval)
            app.invalidate()

    if refresh_interval:
        t = threading.Thread(target=run)
        t.daemon = True
        t.start()

    try:
        yield
    finally:
        # Exit.
        done[0] = True
<|endoftext|>


<commit_before>#!/usr/bin/env python
import dendropy
from dendropy.utility import container
import codecs
import json
import sys
import os
SCRIPT_NAME = os.path.split(sys.argv[0])[1]
class Reason(object):
    NO_INC_DESIGNATORS_IN_TREE = 0
    SUCCESS = 1
    MRCA_HAS_EXCLUDED = 2
    def to_str(c):
        if c == Reason.NO_INC_DESIGNATORS_IN_TREE:
            return 'no specifiers to be included were in the tree'
        if c == Reason.MRCA_HAS_EXCLUDED:
            return 'the include group is paraphyletic with respect to member/members of the exclude group'
        return ''
    to_str = staticmethod(to_str)

def mod_encode_splits(tree, create_dict=True, delete_outdegree_one=True, internal_node_taxa=False):
    """
    Processes splits on a tree, encoding them as bitmask on each edge.
    Adds the following to each edge:
        - `split_bitmask` : a rooted split representation, i.e. a long/bitmask
            where bits corresponding to indices of taxa descended from this
            edge are turned on
    If `create_dict` is True, then the following is added to the tree:
        - `split_edges`:
            [if `tree.is_rooted`]: a dictionary where keys are the
            splits and values are edges.
            [otherwise]: a container.NormalizedBitmaskDictionary where the keys are the
            normalized (unrooted) split representations and the values
            are edges. A normalized split_mask is where the split_bitmask
            is complemented if the right-most bit is not '0' (or just
            the split_bitmask otherwise).
    If `delete_outdegree_one` is True then nodes with one
        will be deleted as they are encountered (this is required
        if the split_edges dictionary is to refer to all edges in the tree).
        Note this will mean that an unrooted tree like '(A,(B,C))' will
        be changed to '(A,B,C)' after this operation!
    """
    taxon_namespace = tree.taxon_namespace
    if taxon_namespace is None:
        taxon_namespace = tree.infer_taxa()
    if create_dict:
        tree.split_edges = {}
        split_map = tree.split_edges
        # if tree.is_rooted:
        #     tree.split_edges = {}
        # else:
        #     atb = taxon_namespace.all_taxa_bitmask()
        #     d = container.NormalizedBitmaskDict(mask=atb)
        #     tree.split_edges = d
        # split_map = tree.split_edges
    if not tree.seed_node:
        return

    if delete_outdegree_one:
        sn = tree.seed_node
        if not tree.is_rooted:
            if len(sn._child_nodes) == 2:
                tree.deroot()
        while len(sn._child_nodes) == 1:
            c = sn._child_nodes[0]
            if len(c._child_nodes) == 0:
                break
            try:
                sn.edge.length += c.edge.length
            except:
                pass
            sn.remove_child(c)
            for gc in c._child_nodes:
                sn.add_child(gc)

    for edge in tree.postorder_edge_iter():
        cm = 0
        h = edge.head_node
        child_nodes = h._child_nodes
        nc = len(child_nodes)
        ocm = None
        if nc > 0:
            if nc == 1 and delete_outdegree_one and edge.tail_node:
                p = edge.tail_node
                assert p
                c = child_nodes[0]
                try:
                    c.edge.length += edge.length
                except:
                    pass
                pos = p._child_nodes.index(h)
                p.insert_child(pos, c)
                p.remove_child(h)
            else:
                for child in child_nodes:
                    cm |= child.edge.split_bitmask
                if internal_node_taxa:
                    t = edge.head_node.taxon
                    if t:
                        ocm = taxon_namespace.taxon_bitmask(t)
                        cm |= ocm
        else:
            t = edge.head_node.taxon
            if t:
                cm = taxon_namespace.taxon_bitmask(t)
        edge.split_bitmask = cm
        if create_dict:
            split_map[cm] = edge
            if ocm is not None:
                split_map[ocm] = edge
    # create normalized bitmasks, where the full (tree) split mask is *not*
    # all the taxa, but only those found on the tree
    if not tree.is_rooted:
        mask = tree.seed_node.edge.split_bitmask
        d = container.NormalizedBitmaskDict(mask=mask)
        for k, v in tree.split_edges.items():
            d[k] = v
        tree.split_edges = d

def debug(msg):
    sys.stderr.write('{s}: {m}\n'.format(s=SCRIPT_NAME, m=msg))

def outcome(target, return_code, dropped_inc, dropped_exc):
    return {'target': target,
            'reason_code': return_code,
            'dropped_inc': dropped_inc,
            'dropped_exc': dropped_exc}
def taxa_in_tree(tree, taxa_list):
    t = []
    dropped = []
    for i in taxa_list:
        ind = tree.label2index.get(i)
        if ind is not None:
            t.append(tree.taxon_namespace[ind])
        else:
            dropped.append(i)
    return t, dropped
def bits_in_tree(tree, taxa_list):
    t, dropped = taxa_in_tree(tree, taxa_list)
    return [tree.label2bit[i.label] for i in t], dropped
def add_node_based_phyloreferenced_annotation(tree, annotation):
    in_tree, dropped_inc = taxa_in_tree(tree, annotation.des)
    if len(in_tree) == 0:
        return outcome(False, Reason.NO_INC_DESIGNATORS_IN_TREE, dropped_inc, None)
    if len(in_tree) == 1:
        mrca = tree.find_node_with_taxon_label(in_tree[0].label)
    else:
        mrca = tree.mrca(taxa=in_tree)
    assert mrca is not None
    return outcome(mrca, Reason.SUCCESS, None, None)

def add_stem_based_phylorefenced_annotation(tree, annotation):
    in_tree, dropped_inc = taxa_in_tree(tree, annotation.des)
    exclude, dropped_ex = bits_in_tree(tree, annotation.exclude_ancs_of)
    if not exclude:
        return outcome(tree.seed_node, Reason.SUCCESS, dropped_inc, dropped_ex)
    exc_bit_set = 0
    for e in exclude:
        exc_bit_set |= e
    if len(in_tree) == 0:
        return outcome(False, Reason.NO_INC_DESIGNATORS_IN_TREE, dropped_inc, dropped_ex)
    if len(in_tree) == 1:
        mrca = tree.find_node_with_taxon_label(in_tree[0].label)
    else:
        mrca = tree.mrca(taxa=in_tree)
    assert mrca is not None
    if mrca.edge.split_bitmask & exc_bit_set:
        return outcome(False, Reason.MRCA_HAS_EXCLUDED, dropped_inc, dropped_ex)
    deepest_valid = mrca
    curr = mrca.parent_node
    while (curr is not None) and ((curr.edge.split_bitmask & exc_bit_set) == 0):
        print 'no intersection of', curr.edge.split_bitmask, exc_bit_set
        deepest_valid = curr
        curr = curr.parent_node
    if curr:
        print 'intersection of', curr.edge.split_bitmask, exc_bit_set
    return outcome(deepest_valid.edge, Reason.SUCCESS, dropped_inc, dropped_ex)

def add_phyloreferenced_annotation(tree, annotation):
    if annotation.rooted_by == GroupType.STEM:
        r = add_stem_based_phylorefenced_annotation(tree, annotation)
    else:
        assert annotation.rooted_by == GroupType.NODE
        r = add_node_based_phyloreferenced_annotation(tree, annotation)
    if r['target']:
        r['target'].phylo_ref.append(annotation)
        annotation.applied_to.append((tree, r['target']))
    return r

class GroupType:
    STEM, NODE = range(2)
    def to_str(c):
        if c == GroupType.STEM:
            return 'stem'
        assert c == GroupType.NODE
        return 'node'
    to_str = staticmethod(to_str)
    def to_code(c):
        if c.lower() == 'stem':
            return GroupType.STEM
        assert c.lower() == 'node'
        return GroupType.NODE
    to_code = staticmethod(to_code)

class PhyloReferencedAnnotation(object):
    def __init__(self, serialized):
        self.des = [str(i) for i in serialized['descendants']]
        self.exclude_ancs_of = [str(i) for i in serialized.get('excludes_ancestors_of', [])]
        self.rooted_by = GroupType.to_code(serialized['rooted_by'])
        self.applied_to = []
    def get_summary(self):
        return json.dumps(self.serialize())
    summary = property(get_summary)
    def serialize(self):
        return {'descendants': self.des,
                'excludes_ancestors_of': self.exclude_ancs_of,
                'rooted_by': GroupType.to_str(self.rooted_by)}
def main(tree_filename, annotations_filename):
    tree_list = dendropy.TreeList.get_from_path(tree_filename,
                                                'newick',
                                                suppress_internal_node_taxa=False)
    for tree in tree_list:
        tree.print_plot(show_internal_node_ids=True)
        mod_encode_splits(tree, delete_outdegree_one=False, internal_node_taxa=True)
        tree.label2index = {}
        tree.label2bit = {}
        curr_bit = 1
        for n, taxon in enumerate(tree.taxon_namespace):
            assert taxon.label not in tree.label2index
            tree.label2index[taxon.label] = n
            tree.label2bit[taxon.label] = curr_bit
            curr_bit <<= 1
        print tree.label2index
        print tree.label2bit
        for node in tree.preorder_node_iter():
            node.phylo_ref = []
            if node.edge:
                node.edge.phylo_ref = []
            print node.edge.split_bitmask
            if node.taxon:
                print node.taxon.label

    a_f = codecs.open(annotations_filename, 'rU', encoding='utf-8')
    annot_list = json.load(a_f)
    if not isinstance(annot_list, list):
        annot_list = [annot_list]

    for tree_index, tree in enumerate(tree_list):
        num_tried = 0
        num_added = 0
        unadded = []
        for annot_index, annotation in enumerate(annot_list):
            a = PhyloReferencedAnnotation(annotation)
            response = add_phyloreferenced_annotation(tree, a)
            x = response['target']
            if x:
                num_added += 1
            else:
                unadded.append((a, response))
                debug('Annotation {a} could not be added to tree {t}'.format(a=annot_index, t=tree_index))
            num_tried += 1
        debug('{a}/{t} annotations added to tree {i}'.format(a=num_added, t=num_tried, i=tree_index))
        # Report tree and annotations
        tree.print_plot(show_internal_node_ids=True)
        for node in tree.preorder_node_iter():
            if node.taxon:
                l = node.taxon.label
            else:
                l = '@' + str(id(node))
            print 'Node "{l}":'.format(l=l)
            if node.phylo_ref:
                for a in node.phylo_ref:
                    print '   ', a.summary
            else:
                print '    <NO ANNOTATIONS>'
            e = node.edge
            if e:
                p = e.tail_node
                if p:
                    if p.taxon:
                        pl = p.taxon.label
                    else:
                        pl = '@' + str(id(p))
                else:
                    pl = 'None'
                print 'Edge "{p}" -> "{l}"'.format(p=pl, l=l)
                if e.phylo_ref:
                    for a in e.phylo_ref:
                        print '   ', a.summary
                else:
                    print '    <NO ANNOTATIONS>'
        # Report unadded annotations
        if len(unadded) > 0:
            print 'Unattached annotations'
        for annotation, add_record in unadded:
            print 'reason={r}. annotation={a}'.format(a=annotation.summary,
                                                      r=Reason.to_str(add_record['reason_code']))

if __name__ == '__main__':
    try:
        tree_file, annotations_file = sys.argv[1:]
    except:
        sys.exit('expecting 2 arguments: a tree file and a JSON file of annotations')
    main(tree_file, annotations_file)
<commit_msg>excludes now working for MRCA queries, in theory<commit_after>#!/usr/bin/env python
import dendropy
from dendropy.utility import container
import codecs
import json
import sys
import os
SCRIPT_NAME = os.path.split(sys.argv[0])[1]
class Reason(object):
    NO_INC_DESIGNATORS_IN_TREE = 0
    SUCCESS = 1
    MRCA_HAS_EXCLUDED = 2
    def to_str(c):
        if c == Reason.NO_INC_DESIGNATORS_IN_TREE:
            return 'no specifiers to be included were in the tree'
        if c == Reason.MRCA_HAS_EXCLUDED:
            return 'the include group is paraphyletic with respect to member/members of the exclude group'
        return ''
    to_str = staticmethod(to_str)

def mod_encode_splits(tree, create_dict=True, delete_outdegree_one=True, internal_node_taxa=False):
    """
    Processes splits on a tree, encoding them as bitmask on each edge.
    Adds the following to each edge:
        - `split_bitmask` : a rooted split representation, i.e. a long/bitmask
            where bits corresponding to indices of taxa descended from this
            edge are turned on
    If `create_dict` is True, then the following is added to the tree:
        - `split_edges`:
            [if `tree.is_rooted`]: a dictionary where keys are the
            splits and values are edges.
            [otherwise]: a container.NormalizedBitmaskDictionary where the keys are the
            normalized (unrooted) split representations and the values
            are edges. A normalized split_mask is where the split_bitmask
            is complemented if the right-most bit is not '0' (or just
            the split_bitmask otherwise).
    If `delete_outdegree_one` is True then nodes with one
        will be deleted as they are encountered (this is required
        if the split_edges dictionary is to refer to all edges in the tree).
        Note this will mean that an unrooted tree like '(A,(B,C))' will
        be changed to '(A,B,C)' after this operation!
    """
    taxon_namespace = tree.taxon_namespace
    if taxon_namespace is None:
        taxon_namespace = tree.infer_taxa()
    if create_dict:
        tree.split_edges = {}
        split_map = tree.split_edges
        # if tree.is_rooted:
        #     tree.split_edges = {}
        # else:
        #     atb = taxon_namespace.all_taxa_bitmask()
        #     d = container.NormalizedBitmaskDict(mask=atb)
        #     tree.split_edges = d
        # split_map = tree.split_edges
    if not tree.seed_node:
        return

    if delete_outdegree_one:
        sn = tree.seed_node
        if not tree.is_rooted:
            if len(sn._child_nodes) == 2:
                tree.deroot()
        while len(sn._child_nodes) == 1:
            c = sn._child_nodes[0]
            if len(c._child_nodes) == 0:
                break
            try:
                sn.edge.length += c.edge.length
            except:
                pass
            sn.remove_child(c)
            for gc in c._child_nodes:
                sn.add_child(gc)

    for edge in tree.postorder_edge_iter():
        cm = 0
        h = edge.head_node
        child_nodes = h._child_nodes
        nc = len(child_nodes)
        ocm = None
        if nc > 0:
            if nc == 1 and delete_outdegree_one and edge.tail_node:
                p = edge.tail_node
                assert p
                c = child_nodes[0]
                try:
                    c.edge.length += edge.length
                except:
                    pass
                pos = p._child_nodes.index(h)
                p.insert_child(pos, c)
                p.remove_child(h)
            else:
                for child in child_nodes:
                    cm |= child.edge.split_bitmask
                if internal_node_taxa:
                    t = edge.head_node.taxon
                    if t:
                        ocm = taxon_namespace.taxon_bitmask(t)
                        cm |= ocm
        else:
            t = edge.head_node.taxon
            if t:
                cm = taxon_namespace.taxon_bitmask(t)
        edge.split_bitmask = cm
        if create_dict:
            split_map[cm] = edge
            if ocm is not None:
                split_map[ocm] = edge
    # create normalized bitmasks, where the full (tree) split mask is *not*
    # all the taxa, but only those found on the tree
    if not tree.is_rooted:
        mask = tree.seed_node.edge.split_bitmask
        d = container.NormalizedBitmaskDict(mask=mask)
        for k, v in tree.split_edges.items():
            d[k] = v
        tree.split_edges = d

def debug(msg):
    sys.stderr.write('{s}: {m}\n'.format(s=SCRIPT_NAME, m=msg))

def outcome(target, return_code, dropped_inc, dropped_exc):
    return {'target': target,
            'reason_code': return_code,
            'dropped_inc': dropped_inc,
            'dropped_exc': dropped_exc}
def taxa_in_tree(tree, taxa_list):
    t = []
    dropped = []
    for i in taxa_list:
        ind = tree.label2index.get(i)
        if ind is not None:
            t.append(tree.taxon_namespace[ind])
        else:
            dropped.append(i)
    return t, dropped
def bits_in_tree(tree, taxa_list):
    t, dropped = taxa_in_tree(tree, taxa_list)
    return [tree.label2bit[i.label] for i in t], dropped
def add_node_based_phyloreferenced_annotation(tree, annotation):
    resp = _find_mrca_and_verify_monophyly(tree, annotation)
    if isinstance(resp, dict):
        return resp
    mrca, exc_bit_set, dropped_inc, dropped_exc = resp
    return outcome(mrca, Reason.SUCCESS, None, None)

def _find_mrca_and_verify_monophyly(tree, annotation):
    in_tree, dropped_inc = taxa_in_tree(tree, annotation.des)
    exclude, dropped_exc = bits_in_tree(tree, annotation.exclude_ancs_of)
    if not exclude:
        return outcome(tree.seed_node, Reason.SUCCESS, dropped_inc, dropped_exc)
    exc_bit_set = 0
    for e in exclude:
        exc_bit_set |= e
    if len(in_tree) == 0:
        return outcome(False, Reason.NO_INC_DESIGNATORS_IN_TREE, dropped_inc, dropped_exc)
    if len(in_tree) == 1:
        mrca = tree.find_node_with_taxon_label(in_tree[0].label)
    else:
        mrca = tree.mrca(taxa=in_tree)
    assert mrca is not None
    if mrca.edge.split_bitmask & exc_bit_set:
        return outcome(False, Reason.MRCA_HAS_EXCLUDED, dropped_inc, dropped_exc)
    return mrca, exc_bit_set, dropped_inc, dropped_exc
def add_stem_based_phylorefenced_annotation(tree, annotation):
    resp = _find_mrca_and_verify_monophyly(tree, annotation)
    if isinstance(resp, dict):
        return resp
    mrca, exc_bit_set, dropped_inc, dropped_exc = resp
    deepest_valid = mrca
    curr = mrca.parent_node
    while (curr is not None) and ((curr.edge.split_bitmask & exc_bit_set) == 0):
        print 'no intersection of', curr.edge.split_bitmask, exc_bit_set
        deepest_valid = curr
        curr = curr.parent_node
    if curr:
        print 'intersection of', curr.edge.split_bitmask, exc_bit_set
    return outcome(deepest_valid.edge, Reason.SUCCESS, dropped_inc, dropped_exc)

def add_phyloreferenced_annotation(tree, annotation):
    if annotation.rooted_by == GroupType.STEM:
        r = add_stem_based_phylorefenced_annotation(tree, annotation)
    else:
        assert annotation.rooted_by == GroupType.NODE
        r = add_node_based_phyloreferenced_annotation(tree, annotation)
    if r['target']:
        r['target'].phylo_ref.append(annotation)
        annotation.applied_to.append((tree, r['target']))
    return r

class GroupType:
    STEM, NODE = range(2)
    def to_str(c):
        if c == GroupType.STEM:
            return 'stem'
        assert c == GroupType.NODE
        return 'node'
    to_str = staticmethod(to_str)
    def to_code(c):
        if c.lower() == 'stem':
            return GroupType.STEM
        assert c.lower() == 'node'
        return GroupType.NODE
    to_code = staticmethod(to_code)

class PhyloReferencedAnnotation(object):
    def __init__(self, serialized):
        self.des = [str(i) for i in serialized['descendants']]
        self.exclude_ancs_of = [str(i) for i in serialized.get('excludes_ancestors_of', [])]
        self.rooted_by = GroupType.to_code(serialized['rooted_by'])
        self.applied_to = []
    def get_summary(self):
        return json.dumps(self.serialize())
    summary = property(get_summary)
    def serialize(self):
        return {'descendants': self.des,
                'excludes_ancestors_of': self.exclude_ancs_of,
                'rooted_by': GroupType.to_str(self.rooted_by)}
def main(tree_filename, annotations_filename):
    tree_list = dendropy.TreeList.get_from_path(tree_filename,
                                                'newick',
                                                suppress_internal_node_taxa=False)
    for tree in tree_list:
        tree.print_plot(show_internal_node_ids=True)
        mod_encode_splits(tree, delete_outdegree_one=False, internal_node_taxa=True)
        tree.label2index = {}
        tree.label2bit = {}
        curr_bit = 1
        for n, taxon in enumerate(tree.taxon_namespace):
            assert taxon.label not in tree.label2index
            tree.label2index[taxon.label] = n
            tree.label2bit[taxon.label] = curr_bit
            curr_bit <<= 1
        print tree.label2index
        print tree.label2bit
        for node in tree.preorder_node_iter():
            node.phylo_ref = []
            if node.edge:
                node.edge.phylo_ref = []
            print node.edge.split_bitmask
            if node.taxon:
                print node.taxon.label

    a_f = codecs.open(annotations_filename, 'rU', encoding='utf-8')
    annot_list = json.load(a_f)
    if not isinstance(annot_list, list):
        annot_list = [annot_list]

    for tree_index, tree in enumerate(tree_list):
        num_tried = 0
        num_added = 0
        unadded = []
        for annot_index, annotation in enumerate(annot_list):
            a = PhyloReferencedAnnotation(annotation)
            response = add_phyloreferenced_annotation(tree, a)
            x = response['target']
            if x:
                num_added += 1
            else:
                unadded.append((a, response))
                debug('Annotation {a} could not be added to tree {t}'.format(a=annot_index, t=tree_index))
            num_tried += 1
        debug('{a}/{t} annotations added to tree {i}'.format(a=num_added, t=num_tried, i=tree_index))
        # Report tree and annotations
        tree.print_plot(show_internal_node_ids=True)
        for node in tree.preorder_node_iter():
            if node.taxon:
                l = node.taxon.label
            else:
                l = '@' + str(id(node))
            print 'Node "{l}":'.format(l=l)
            if node.phylo_ref:
                for a in node.phylo_ref:
                    print '   ', a.summary
            else:
                print '    <NO ANNOTATIONS>'
            e = node.edge
            if e:
                p = e.tail_node
                if p:
                    if p.taxon:
                        pl = p.taxon.label
                    else:
                        pl = '@' + str(id(p))
                else:
                    pl = 'None'
                print 'Edge "{p}" -> "{l}"'.format(p=pl, l=l)
                if e.phylo_ref:
                    for a in e.phylo_ref:
                        print '   ', a.summary
                else:
                    print '    <NO ANNOTATIONS>'
        # Report unadded annotations
        if len(unadded) > 0:
            print 'Unattached annotations'
        for annotation, add_record in unadded:
            print 'reason={r}. annotation={a}'.format(a=annotation.summary,
                                                      r=Reason.to_str(add_record['reason_code']))

if __name__ == '__main__':
    try:
        tree_file, annotations_file = sys.argv[1:]
    except:
        sys.exit('expecting 2 arguments: a tree file and a JSON file of annotations')
    main(tree_file, annotations_file)
<|endoftext|>


<commit_before># django imports
from django.core.urlresolvers import reverse
from django.db import models

# third party imports
from db_file_storage.model_utils import delete_file, delete_file_if_needed

class CDDisc(models.Model):
    bytes = models.TextField()
    filename = models.CharField(max_length=255)
    mimetype = models.CharField(max_length=50)

class CDCover(models.Model):
    bytes = models.TextField()
    filename = models.CharField(max_length=255)
    mimetype = models.CharField(max_length=50)

class CD(models.Model):
    name = models.CharField(max_length=100, unique=True)
    disc = models.ImageField(
        upload_to='music.CDDisc/bytes/filename/mimetype',
        blank=True,
        null=True
    )
    cover = models.ImageField(
        upload_to='music.CDCover/bytes/filename/mimetype',
        blank=True,
        null=True
    )

    def __unicode__(self):
        return self.name

    def get_absolute_url(self):
        return reverse('cd.edit', kwargs={'pk': self.pk})

    def save(self, *args, **kwargs):
        delete_file_if_needed(self, 'disc')
        delete_file_if_needed(self, 'cover')
        super(CD, self).save(*args, **kwargs)

    def delete(self, *args, **kwargs):
        super(CD, self).delete(*args, **kwargs)
        delete_file(self, 'disc')
        delete_file(self, 'cover')

class SoundDeviceInstructionManual(models.Model):
    bytes = models.TextField()
    filename = models.CharField(max_length=255)
    mimetype = models.CharField(max_length=50)

class SoundDevice(models.Model):
    name = models.CharField(max_length=100)
    instruction_manual = models.FileField(
        upload_to='music.SoundDeviceInstructionManual/bytes/filename/mimetype',
        blank=True,
        null=True
    )
<commit_msg>fixing code smells<commit_after># django imports
from django.core.urlresolvers import reverse
from django.db import models

# third party imports
from db_file_storage.model_utils import delete_file, delete_file_if_needed

class CDDisc(models.Model):
    bytes = models.TextField()
    filename = models.CharField(max_length=255)
    mimetype = models.CharField(max_length=50)

    def __unicode__(self):
        return self.filename

class CDCover(models.Model):
    bytes = models.TextField()
    filename = models.CharField(max_length=255)
    mimetype = models.CharField(max_length=50)

    def __unicode__(self):
        return self.filename

class CD(models.Model):
    name = models.CharField(max_length=100, unique=True)
    disc = models.ImageField(
        upload_to='music.CDDisc/bytes/filename/mimetype',
        blank=True,
        null=True
    )
    cover = models.ImageField(
        upload_to='music.CDCover/bytes/filename/mimetype',
        blank=True,
        null=True
    )

    def __unicode__(self):
        return self.name

    def get_absolute_url(self):
        return reverse('cd.edit', kwargs={'pk': self.pk})

    def save(self, *args, **kwargs):
        delete_file_if_needed(self, 'disc')
        delete_file_if_needed(self, 'cover')
        super(CD, self).save(*args, **kwargs)

    def delete(self, *args, **kwargs):
        super(CD, self).delete(*args, **kwargs)
        delete_file(self, 'disc')
        delete_file(self, 'cover')

class SoundDeviceInstructionManual(models.Model):
    bytes = models.TextField()
    filename = models.CharField(max_length=255)
    mimetype = models.CharField(max_length=50)

    def __unicode__(self):
        return self.filename

class SoundDevice(models.Model):
    name = models.CharField(max_length=100)
    instruction_manual = models.FileField(
        upload_to='music.SoundDeviceInstructionManual/bytes/filename/mimetype',
        blank=True,
        null=True
    )

    def __unicode__(self):
        return self.name
<|endoftext|>


<commit_before>import itertools
import unittest

import numpy
import six

import chainer
from chainer.backends import cuda
from chainer import functions
from chainer import testing
from chainer.testing import attr
from chainer.utils import force_array
from chainer.utils import type_check

def _testing_pairwise(*parameters):
    rs = numpy.random.RandomState(seed=0)
    parameters = [list(enumerate(dicts)) for dicts in parameters]

    uncovered = {}
    for (i, dicts_i), (j, dicts_j) in itertools.combinations(
            enumerate(parameters), 2):
        uncovered[(i, j)] = set(itertools.product(
            range(len(dicts_i)), range(len(dicts_j))))

    product_dicts = list(itertools.product(*parameters))
    rs.shuffle(product_dicts)
    for product_dict in product_dicts:
        count = 0
        ks, dicts = zip(*product_dict)
        for (i, k_i), (j, k_j) in itertools.combinations(
                enumerate(ks), 2):
            # uncovered[(i, j)].discard((k_i, k_j))
            try:
                uncovered[(i, j)].remove((k_i, k_j))
            except KeyError:
                pass
            else:
                count += 1
        if count > 0:
            yield {k: v for dic in dicts for k, v in dic.items()}

def accuracy(x, t, ignore_label):
    x_ = numpy.rollaxis(x, 1, x.ndim).reshape(t.size, -1)
    t_ = t.ravel()

    if ignore_label is not None:
        count = 0
        for i in six.moves.range(t_.size):
            pred = x_[i].argmax()
            if t_[i] != ignore_label and pred == t_[i]:
                count += 1
        total = (t_ != ignore_label).sum()
    else:
        count = 0
        for i in six.moves.range(t_.size):
            pred = x_[i].argmax()
            if pred == t_[i]:
                count += 1
        total = t_.size

    if total == 0:
        return 0.0
    else:
        return float(count) / total

@testing.parameterize(
    *_testing_pairwise(
        [{'x_shape': (10, 3), 't_shape': (10,)},
         {'x_shape': (10, 3, 1), 't_shape': (10,)},
         {'x_shape': (10, 3, 1, 1), 't_shape': (10,)},
         {'x_shape': (10, 3, 5), 't_shape': (10, 5)},
         {'x_shape': (10, 3, 5, 4), 't_shape': (10, 5, 4)},
         {'x_shape': (10, 3, 5, 4, 1), 't_shape': (10, 5, 4)},
         {'x_shape': (10, 3, 5, 4, 1, 1), 't_shape': (10, 5, 4)}],
        [{'ignore_label': None, 't_data': 'randint'},
         {'ignore_label': 0, 't_data': 'randint'},
         {'ignore_label': 0, 't_data': 'zero'}],
        [{'dtype': numpy.float16},
         {'dtype': numpy.float32},
         {'dtype': numpy.float64}],
        [{'label_dtype': numpy.int8},
         {'label_dtype': numpy.int16},
         {'label_dtype': numpy.int32},
         {'label_dtype': numpy.int64}]
    )
)
@testing.fix_random()
@testing.inject_backend_tests(
    None,
    # CPU tests
    [
        {},
    ]
    # GPU tests
    + testing.product({
        'use_cuda': [True],
        'cuda_device': [0, 1],
    })
    # ChainerX tests
    + testing.product({
        'use_chainerx': [True],
        'chainerx_device': ['native:0', 'cuda:0', 'cuda:1'],
    })
)
class TestAccuracy(testing.FunctionTestCase):

    def setUp(self):
        self.skip_backward_test = True
        self.skip_double_backward_test = True

        if self.dtype == numpy.float16:
            self.check_forward_options.update({'atol': 1e-4, 'rtol': 1e-3})

    def generate_inputs(self):
        x = numpy.random.uniform(-1, 1, self.x_shape).astype(self.dtype)
        if self.t_data == 'randint':
            t = numpy.random.randint(
                3, size=self.t_shape).astype(self.label_dtype)
        elif self.t_data == 'zero':
            t = numpy.zeros(self.t_shape).astype(self.label_dtype)
        return x, t

    def forward(self, inputs, device):
        x, t = inputs
        return functions.accuracy(x, t, self.ignore_label),

    def forward_expected(self, inputs):
        x, t = inputs
        expected = accuracy(x, t, self.ignore_label)
        expected = force_array(expected, self.dtype)
        return expected,

@testing.parameterize(
    {'x_shape': (10, 3), 't_shape': (4,)},
    {'x_shape': (10, 3, 2), 't_shape': (10,)},
    {'x_shape': (10, 3, 1, 2), 't_shape': (10,)},
    {'x_shape': (10, 3, 4), 't_shape': (10, 5)},
    {'x_shape': (10, 3, 5, 2), 't_shape': (10, 5)},
    {'x_shape': (10, 3, 5, 1, 2), 't_shape': (10, 5)},
)
class TestInvalidShape(unittest.TestCase):

    def setUp(self):
        self.x = numpy.random.uniform(-1, 1,
                                      self.x_shape).astype(numpy.float32)
        self.t = numpy.random.randint(3, size=self.t_shape).astype(numpy.int32)

    def check_invalid_shape(self, xp):
        x = chainer.Variable(xp.asarray(self.x))
        t = chainer.Variable(xp.asarray(self.t))
        with self.assertRaises(type_check.InvalidType):
            functions.accuracy(x, t)

    def test_invalid_shape_cpu(self):
        self.check_invalid_shape(numpy)

    @attr.gpu
    def test_invalid_shape_gpu(self):
        self.check_invalid_shape(cuda.cupy)

testing.run_module(__name__, __file__)
<commit_msg>Add comment<commit_after>import itertools
import unittest

import numpy
import six

import chainer
from chainer.backends import cuda
from chainer import functions
from chainer import testing
from chainer.testing import attr
from chainer.utils import force_array
from chainer.utils import type_check

def _testing_pairwise(*parameters):
    """Generate combinations that cover all pairs.

    The argument is the same as `chainer.testing.product_dict`.

    """
    # parameters: [[dict<str, _>]]
    rs = numpy.random.RandomState(seed=0)
    # parameters: [[(int, dict)]]
    parameters = [list(enumerate(dicts)) for dicts in parameters]

    # uncovered: dict<(int, int), set<(int, int)>>
    # `(k_i, k_j) in uncovered[(i, j)]` iff there has not been a combination
    # that selects k_i-th item of i-th choice and k_j-th item of j-th choice.
    uncovered = {}
    for (i, dicts_i), (j, dicts_j) in itertools.combinations(
            enumerate(parameters), 2):
        uncovered[(i, j)] = set(itertools.product(
            range(len(dicts_i)), range(len(dicts_j))))

    # product_dicts: [[(int, dict)]]
    product_dicts = list(itertools.product(*parameters))
    rs.shuffle(product_dicts)
    for product_dict in product_dicts:
        count = 0
        # ks: [int]
        # dicts: [dict]
        ks, dicts = zip(*product_dict)
        for (i, k_i), (j, k_j) in itertools.combinations(
                enumerate(ks), 2):
            # uncovered[(i, j)].discard((k_i, k_j))
            try:
                uncovered[(i, j)].remove((k_i, k_j))
            except KeyError:
                pass
            else:
                count += 1
        if count > 0:
            yield {k: v for dic in dicts for k, v in dic.items()}

def accuracy(x, t, ignore_label):
    x_ = numpy.rollaxis(x, 1, x.ndim).reshape(t.size, -1)
    t_ = t.ravel()

    if ignore_label is not None:
        count = 0
        for i in six.moves.range(t_.size):
            pred = x_[i].argmax()
            if t_[i] != ignore_label and pred == t_[i]:
                count += 1
        total = (t_ != ignore_label).sum()
    else:
        count = 0
        for i in six.moves.range(t_.size):
            pred = x_[i].argmax()
            if pred == t_[i]:
                count += 1
        total = t_.size

    if total == 0:
        return 0.0
    else:
        return float(count) / total

@testing.parameterize(
    *_testing_pairwise(
        [{'x_shape': (10, 3), 't_shape': (10,)},
         {'x_shape': (10, 3, 1), 't_shape': (10,)},
         {'x_shape': (10, 3, 1, 1), 't_shape': (10,)},
         {'x_shape': (10, 3, 5), 't_shape': (10, 5)},
         {'x_shape': (10, 3, 5, 4), 't_shape': (10, 5, 4)},
         {'x_shape': (10, 3, 5, 4, 1), 't_shape': (10, 5, 4)},
         {'x_shape': (10, 3, 5, 4, 1, 1), 't_shape': (10, 5, 4)}],
        [{'ignore_label': None, 't_data': 'randint'},
         {'ignore_label': 0, 't_data': 'randint'},
         {'ignore_label': 0, 't_data': 'zero'}],
        [{'dtype': numpy.float16},
         {'dtype': numpy.float32},
         {'dtype': numpy.float64}],
        [{'label_dtype': numpy.int8},
         {'label_dtype': numpy.int16},
         {'label_dtype': numpy.int32},
         {'label_dtype': numpy.int64}]
    )
)
@testing.fix_random()
@testing.inject_backend_tests(
    None,
    # CPU tests
    [
        {},
    ]
    # GPU tests
    + testing.product({
        'use_cuda': [True],
        'cuda_device': [0, 1],
    })
    # ChainerX tests
    + testing.product({
        'use_chainerx': [True],
        'chainerx_device': ['native:0', 'cuda:0', 'cuda:1'],
    })
)
class TestAccuracy(testing.FunctionTestCase):

    def setUp(self):
        self.skip_backward_test = True
        self.skip_double_backward_test = True

        if self.dtype == numpy.float16:
            self.check_forward_options.update({'atol': 1e-4, 'rtol': 1e-3})

    def generate_inputs(self):
        x = numpy.random.uniform(-1, 1, self.x_shape).astype(self.dtype)
        if self.t_data == 'randint':
            t = numpy.random.randint(
                3, size=self.t_shape).astype(self.label_dtype)
        elif self.t_data == 'zero':
            t = numpy.zeros(self.t_shape).astype(self.label_dtype)
        return x, t

    def forward(self, inputs, device):
        x, t = inputs
        return functions.accuracy(x, t, self.ignore_label),

    def forward_expected(self, inputs):
        x, t = inputs
        expected = accuracy(x, t, self.ignore_label)
        expected = force_array(expected, self.dtype)
        return expected,

@testing.parameterize(
    {'x_shape': (10, 3), 't_shape': (4,)},
    {'x_shape': (10, 3, 2), 't_shape': (10,)},
    {'x_shape': (10, 3, 1, 2), 't_shape': (10,)},
    {'x_shape': (10, 3, 4), 't_shape': (10, 5)},
    {'x_shape': (10, 3, 5, 2), 't_shape': (10, 5)},
    {'x_shape': (10, 3, 5, 1, 2), 't_shape': (10, 5)},
)
class TestInvalidShape(unittest.TestCase):

    def setUp(self):
        self.x = numpy.random.uniform(-1, 1,
                                      self.x_shape).astype(numpy.float32)
        self.t = numpy.random.randint(3, size=self.t_shape).astype(numpy.int32)

    def check_invalid_shape(self, xp):
        x = chainer.Variable(xp.asarray(self.x))
        t = chainer.Variable(xp.asarray(self.t))
        with self.assertRaises(type_check.InvalidType):
            functions.accuracy(x, t)

    def test_invalid_shape_cpu(self):
        self.check_invalid_shape(numpy)

    @attr.gpu
    def test_invalid_shape_gpu(self):
        self.check_invalid_shape(cuda.cupy)

testing.run_module(__name__, __file__)
<|endoftext|>


<commit_before>from tests.unit.dataactcore.factories.domain import TASFactory
from tests.unit.dataactcore.factories.staging import (
    AwardFinancialFactory, ObjectClassProgramActivityFactory)
from tests.unit.dataactvalidator.utils import number_of_errors, query_columns

_FILE = 'b20_object_class_program_activity'

def test_column_headers(database):
    expected_subset = {'row_number', 'allocation_transfer_agency', 'agency_identifier',
                       'beginning_period_of_availa', 'ending_period_of_availabil', 'availability_type_code',
                       'main_account_code', 'sub_account_code', 'program_activity_code', 'object_class'}
    actual = set(query_columns(_FILE, database))
    assert (actual & expected_subset) == expected_subset

def test_success(database):
    """ Tests that all combinations of TAS, program activity code, and object class in File C exist in File B """
    tas = TASFactory()
    database.session.add(tas)
    database.session.flush()

    op = ObjectClassProgramActivityFactory(tas_id=tas.tas_id, program_activity_code='1', object_class='1')
    af = AwardFinancialFactory(tas_id=tas.tas_id, program_activity_code='1', object_class='1')
    # Allow program activity code to be null, empty, or zero
    af2 = AwardFinancialFactory(tas_id=tas.tas_id, program_activity_code='', object_class='1')
    af3 = AwardFinancialFactory(tas_id=tas.tas_id, program_activity_code='0000', object_class='1')
    af4 = AwardFinancialFactory(tas_id=tas.tas_id, program_activity_code='0', object_class='1')
    af5 = AwardFinancialFactory(tas_id=tas.tas_id, program_activity_code=None, object_class='1')

    assert number_of_errors(_FILE, database, models=[op, af, af2, af3, af4, af5]) == 0

def test_failure(database):
    """ Tests that all combinations of TAS, program activity code, and object class in File C do not exist in File B """
    tas1 = TASFactory()
    tas2 = TASFactory()
    database.session.add_all([tas1, tas2])
    database.session.flush()

    op = ObjectClassProgramActivityFactory(tas_id=tas1.tas_id, program_activity_code='1', object_class='1')

    af1 = AwardFinancialFactory(tas_id=tas1.tas_id, program_activity_code='1', object_class='1')
    af2 = AwardFinancialFactory(tas_id=tas2.tas_id, program_activity_code='1', object_class='1')
    af3 = AwardFinancialFactory(tas_id=tas1.tas_id, program_activity_code='2', object_class='1')
    af4 = AwardFinancialFactory(tas_id=tas1.tas_id, program_activity_code='1', object_class='2')

    assert number_of_errors(_FILE, database, models=[op, af1, af2, af3, af4]) == 3
<commit_msg>Revert code back to original<commit_after>from tests.unit.dataactcore.factories.domain import TASFactory
from tests.unit.dataactcore.factories.staging import (
    AwardFinancialFactory, ObjectClassProgramActivityFactory)
from tests.unit.dataactvalidator.utils import number_of_errors, query_columns

_FILE = 'b20_object_class_program_activity'

def test_column_headers(database):
    expected_subset = {'row_number', 'allocation_transfer_agency', 'agency_identifier',
                       'beginning_period_of_availa', 'ending_period_of_availabil', 'availability_type_code',
                       'main_account_code', 'sub_account_code', 'program_activity_code', 'object_class'}
    actual = set(query_columns(_FILE, database))
    assert (actual & expected_subset) == expected_subset

def test_success(database):
    """ Tests that all combinations of TAS, program activity code, and object class in File C exist in File B """
    tas = TASFactory()
    database.session.add(tas)
    database.session.flush()

    op = ObjectClassProgramActivityFactory(tas_id=tas.tas_id, program_activity_code='1', object_class='1')
    af = AwardFinancialFactory(tas_id=tas.tas_id, program_activity_code='1', object_class='1')
    # Allow program activity code to be null, empty, or zero
    af2 = AwardFinancialFactory(tas_id=tas.tas_id, program_activity_code='', object_class='1')
    af3 = AwardFinancialFactory(tas_id=tas.tas_id, program_activity_code='0000', object_class='1')
    af4 = AwardFinancialFactory(tas_id=tas.tas_id, program_activity_code=None, object_class='1')

    assert number_of_errors(_FILE, database, models=[op, af, af2, af3, af4]) == 0

def test_failure(database):
    """ Tests that all combinations of TAS, program activity code, and object class in File C do not exist in File B """
    tas1 = TASFactory()
    tas2 = TASFactory()
    database.session.add_all([tas1, tas2])
    database.session.flush()

    op = ObjectClassProgramActivityFactory(tas_id=tas1.tas_id, program_activity_code='1', object_class='1')

    af1 = AwardFinancialFactory(tas_id=tas1.tas_id, program_activity_code='1', object_class='1')
    af2 = AwardFinancialFactory(tas_id=tas2.tas_id, program_activity_code='1', object_class='1')
    af3 = AwardFinancialFactory(tas_id=tas1.tas_id, program_activity_code='2', object_class='1')
    af4 = AwardFinancialFactory(tas_id=tas1.tas_id, program_activity_code='1', object_class='2')

    assert number_of_errors(_FILE, database, models=[op, af1, af2, af3, af4]) == 3
<|endoftext|>


<commit_before># coding=utf-8
# Copyright 2020 The TF-Agents Authors.
#
# 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
#
#     https://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.

"""Runs training and eval on CQL-SAC on D4RL using the Actor-Learner API.

All default hyperparameters in train_eval come from the CQL paper:
https://arxiv.org/abs/2006.04779
"""
import os

from absl import app
from absl import flags
from absl import logging

import gin
import numpy as np
import tensorflow as tf

from tf_agents.agents.cql import cql_sac_agent
from tf_agents.agents.ddpg import critic_network
from tf_agents.agents.sac import tanh_normal_projection_network
from tf_agents.environments import tf_py_environment
from tf_agents.experimental.examples.cql_sac.kumar20.d4rl_utils import load_d4rl
from tf_agents.experimental.examples.cql_sac.kumar20.data_utils import create_tf_record_dataset
from tf_agents.metrics import py_metrics
from tf_agents.networks import actor_distribution_network
from tf_agents.policies import py_tf_eager_policy
from tf_agents.train import actor
from tf_agents.train import learner
from tf_agents.train import triggers
from tf_agents.train.utils import strategy_utils
from tf_agents.train.utils import train_utils
from tf_agents.trajectories import trajectory

FLAGS = flags.FLAGS

flags.DEFINE_string('root_dir', os.getenv('TEST_UNDECLARED_OUTPUTS_DIR'),
                    'Root directory for writing logs/summaries/checkpoints.')
flags.DEFINE_string('env_name', 'antmaze-medium-play-v0',
                    'Name of the environment.')
flags.DEFINE_string('dataset_path', None, 'TFRecord dataset path.')
flags.DEFINE_integer('learner_iterations_per_call', 500,
                     'Iterations per learner run call.')
flags.DEFINE_integer('policy_save_interval', 10000, 'Policy save interval.')
flags.DEFINE_integer('eval_interval', 10000, 'Evaluation interval.')
flags.DEFINE_integer('summary_interval', 1000, 'Summary interval.')
flags.DEFINE_integer('num_gradient_updates', 1000000,
                     'Total number of train iterations to perform.')
flags.DEFINE_bool(
    'use_trajectories', False,
    'Whether dataset samples are stored as trajectories. '
    'If False, stored as transitions')
flags.DEFINE_multi_string('gin_file', None, 'Paths to the gin-config files.')
flags.DEFINE_multi_string('gin_param', None, 'Gin binding parameters.')

@gin.configurable
def train_eval(
    root_dir,
    dataset_path,
    env_name,
    # Training params
    tpu=False,
    use_gpu=False,
    num_gradient_updates=1000000,
    actor_fc_layers=(256, 256),
    critic_joint_fc_layers=(256, 256, 256),
    # Agent params
    batch_size=256,
    bc_steps=0,
    actor_learning_rate=3e-5,
    critic_learning_rate=3e-4,
    alpha_learning_rate=3e-4,
    reward_scale_factor=1.0,
    cql_alpha_learning_rate=3e-4,
    cql_alpha=5.0,
    cql_tau=10.0,
    num_cql_samples=10,
    reward_noise_variance=0.0,
    include_critic_entropy_term=False,
    use_lagrange_cql_alpha=True,
    log_cql_alpha_clipping=None,
    softmax_temperature=1.0,
    # Data params
    reward_shift=0.0,
    action_clipping=None,
    use_trajectories=False,
    data_shuffle_buffer_size_per_record=1,
    data_shuffle_buffer_size=100,
    data_num_shards=1,
    data_block_length=10,
    data_parallel_reads=None,
    data_parallel_calls=10,
    data_prefetch=10,
    data_cycle_length=10,
    # Others
    policy_save_interval=10000,
    eval_interval=10000,
    summary_interval=1000,
    learner_iterations_per_call=1,
    eval_episodes=10,
    debug_summaries=False,
    summarize_grads_and_vars=False,
    seed=None):
  """Trains and evaluates CQL-SAC."""
  logging.info('Training CQL-SAC on: %s', env_name)
  tf.random.set_seed(seed)
  np.random.seed(seed)

  # Load environment.
  env = load_d4rl(env_name)
  tf_env = tf_py_environment.TFPyEnvironment(env)
  strategy = strategy_utils.get_strategy(tpu, use_gpu)

  if not dataset_path.endswith('.tfrecord'):
    dataset_path = os.path.join(dataset_path, env_name,
                                '%s*.tfrecord' % env_name)
  logging.info('Loading dataset from %s', dataset_path)
  dataset_paths = tf.io.gfile.glob(dataset_path)

  # Create dataset.
  with strategy.scope():
    dataset = create_tf_record_dataset(
        dataset_paths,
        batch_size,
        shuffle_buffer_size_per_record=data_shuffle_buffer_size_per_record,
        shuffle_buffer_size=data_shuffle_buffer_size,
        num_shards=data_num_shards,
        cycle_length=data_cycle_length,
        block_length=data_block_length,
        num_parallel_reads=data_parallel_reads,
        num_parallel_calls=data_parallel_calls,
        num_prefetch=data_prefetch,
        strategy=strategy,
        reward_shift=reward_shift,
        action_clipping=action_clipping,
        use_trajectories=use_trajectories)

  # Create agent.
  time_step_spec = tf_env.time_step_spec()
  observation_spec = time_step_spec.observation
  action_spec = tf_env.action_spec()
  with strategy.scope():
    train_step = train_utils.create_train_step()

    actor_net = actor_distribution_network.ActorDistributionNetwork(
        observation_spec,
        action_spec,
        fc_layer_params=actor_fc_layers,
        continuous_projection_net=tanh_normal_projection_network
        .TanhNormalProjectionNetwork)

    critic_net = critic_network.CriticNetwork(
        (observation_spec, action_spec),
        joint_fc_layer_params=critic_joint_fc_layers,
        kernel_initializer='glorot_uniform',
        last_kernel_initializer='glorot_uniform')

    agent = cql_sac_agent.CqlSacAgent(
        time_step_spec,
        action_spec,
        actor_network=actor_net,
        critic_network=critic_net,
        actor_optimizer=tf.keras.optimizers.Adam(
            learning_rate=actor_learning_rate),
        critic_optimizer=tf.keras.optimizers.Adam(
            learning_rate=critic_learning_rate),
        alpha_optimizer=tf.keras.optimizers.Adam(
            learning_rate=alpha_learning_rate),
        cql_alpha=cql_alpha,
        num_cql_samples=num_cql_samples,
        include_critic_entropy_term=include_critic_entropy_term,
        use_lagrange_cql_alpha=use_lagrange_cql_alpha,
        cql_alpha_learning_rate=cql_alpha_learning_rate,
        target_update_tau=5e-3,
        target_update_period=1,
        random_seed=seed,
        cql_tau=cql_tau,
        reward_noise_variance=reward_noise_variance,
        num_bc_steps=bc_steps,
        td_errors_loss_fn=tf.math.squared_difference,
        gamma=0.99,
        reward_scale_factor=reward_scale_factor,
        gradient_clipping=None,
        log_cql_alpha_clipping=log_cql_alpha_clipping,
        softmax_temperature=softmax_temperature,
        debug_summaries=debug_summaries,
        summarize_grads_and_vars=summarize_grads_and_vars,
        train_step_counter=train_step)
    agent.initialize()

  # Create learner.
  saved_model_dir = os.path.join(root_dir, learner.POLICY_SAVED_MODEL_DIR)
  collect_env_step_metric = py_metrics.EnvironmentSteps()
  learning_triggers = [
      triggers.PolicySavedModelTrigger(
          saved_model_dir,
          agent,
          train_step,
          interval=policy_save_interval,
          metadata_metrics={
              triggers.ENV_STEP_METADATA_KEY: collect_env_step_metric
          }),
      triggers.StepPerSecondLogTrigger(train_step, interval=100)
  ]
  cql_learner = learner.Learner(
      root_dir,
      train_step,
      agent,
      experience_dataset_fn=lambda: dataset,
      triggers=learning_triggers,
      summary_interval=summary_interval,
      strategy=strategy)

  # Create actor for evaluation.
  eval_greedy_policy = py_tf_eager_policy.PyTFEagerPolicy(
      agent.policy, use_tf_function=True)
  eval_actor = actor.Actor(
      env,
      eval_greedy_policy,
      train_step,
      metrics=actor.eval_metrics(eval_episodes),
      summary_dir=os.path.join(root_dir, 'eval'),
      episodes_per_run=eval_episodes)

  # Run.
  dummy_trajectory = trajectory.mid((), (), (), 0., 1.)
  num_learner_iterations = int(num_gradient_updates /
                               learner_iterations_per_call)
  for _ in range(num_learner_iterations):
    # Mimic collecting environment steps since we loaded a static dataset.
    for _ in range(learner_iterations_per_call):
      collect_env_step_metric(dummy_trajectory)

    cql_learner.run(iterations=learner_iterations_per_call)
    if eval_interval and train_step.numpy() % eval_interval == 0:
      eval_actor.run_and_log()

def main(_):
  logging.set_verbosity(logging.INFO)

  gin.parse_config_files_and_bindings(FLAGS.gin_file, FLAGS.gin_param)

  train_eval(
      root_dir=FLAGS.root_dir,
      dataset_path=FLAGS.dataset_path,
      env_name=FLAGS.env_name,
      tpu=FLAGS.tpu,
      use_gpu=FLAGS.use_gpu,
      num_gradient_updates=FLAGS.num_gradient_updates,
      policy_save_interval=FLAGS.policy_save_interval,
      eval_interval=FLAGS.eval_interval,
      summary_interval=FLAGS.summary_interval,
      learner_iterations_per_call=FLAGS.learner_iterations_per_call,
      use_trajectories=FLAGS.use_trajectories)

if __name__ == '__main__':
  flags.mark_flag_as_required('root_dir')
  flags.mark_flag_as_required('dataset_path')
  app.run(main)
<commit_msg>Changes eval to use greedy_policy improving half-cheetah-medium-v0 from ~3.2K to ~4.8K average return during 1M step run.<commit_after># coding=utf-8
# Copyright 2020 The TF-Agents Authors.
#
# 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
#
#     https://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.

"""Runs training and eval on CQL-SAC on D4RL using the Actor-Learner API.

All default hyperparameters in train_eval come from the CQL paper:
https://arxiv.org/abs/2006.04779
"""
import os

from absl import app
from absl import flags
from absl import logging

import gin
import numpy as np
import tensorflow as tf

from tf_agents.agents.cql import cql_sac_agent
from tf_agents.agents.ddpg import critic_network
from tf_agents.agents.sac import tanh_normal_projection_network
from tf_agents.environments import tf_py_environment
from tf_agents.experimental.examples.cql_sac.kumar20.d4rl_utils import load_d4rl
from tf_agents.experimental.examples.cql_sac.kumar20.data_utils import create_tf_record_dataset
from tf_agents.metrics import py_metrics
from tf_agents.networks import actor_distribution_network
from tf_agents.policies import greedy_policy
from tf_agents.policies import py_tf_eager_policy
from tf_agents.train import actor
from tf_agents.train import learner
from tf_agents.train import triggers
from tf_agents.train.utils import strategy_utils
from tf_agents.train.utils import train_utils
from tf_agents.trajectories import trajectory

FLAGS = flags.FLAGS

flags.DEFINE_string('root_dir', os.getenv('TEST_UNDECLARED_OUTPUTS_DIR'),
                    'Root directory for writing logs/summaries/checkpoints.')
flags.DEFINE_string('env_name', 'antmaze-medium-play-v0',
                    'Name of the environment.')
flags.DEFINE_string('dataset_path', None, 'TFRecord dataset path.')
flags.DEFINE_integer('learner_iterations_per_call', 500,
                     'Iterations per learner run call.')
flags.DEFINE_integer('policy_save_interval', 10000, 'Policy save interval.')
flags.DEFINE_integer('eval_interval', 10000, 'Evaluation interval.')
flags.DEFINE_integer('summary_interval', 1000, 'Summary interval.')
flags.DEFINE_integer('num_gradient_updates', 1000000,
                     'Total number of train iterations to perform.')
flags.DEFINE_bool(
    'use_trajectories', False,
    'Whether dataset samples are stored as trajectories. '
    'If False, stored as transitions')
flags.DEFINE_multi_string('gin_file', None, 'Paths to the gin-config files.')
flags.DEFINE_multi_string('gin_param', None, 'Gin binding parameters.')

@gin.configurable
def train_eval(
    root_dir,
    dataset_path,
    env_name,
    # Training params
    tpu=False,
    use_gpu=False,
    num_gradient_updates=1000000,
    actor_fc_layers=(256, 256),
    critic_joint_fc_layers=(256, 256, 256),
    # Agent params
    batch_size=256,
    bc_steps=0,
    actor_learning_rate=3e-5,
    critic_learning_rate=3e-4,
    alpha_learning_rate=3e-4,
    reward_scale_factor=1.0,
    cql_alpha_learning_rate=3e-4,
    cql_alpha=5.0,
    cql_tau=10.0,
    num_cql_samples=10,
    reward_noise_variance=0.0,
    include_critic_entropy_term=False,
    use_lagrange_cql_alpha=True,
    log_cql_alpha_clipping=None,
    softmax_temperature=1.0,
    # Data params
    reward_shift=0.0,
    action_clipping=None,
    use_trajectories=False,
    data_shuffle_buffer_size_per_record=1,
    data_shuffle_buffer_size=100,
    data_num_shards=1,
    data_block_length=10,
    data_parallel_reads=None,
    data_parallel_calls=10,
    data_prefetch=10,
    data_cycle_length=10,
    # Others
    policy_save_interval=10000,
    eval_interval=10000,
    summary_interval=1000,
    learner_iterations_per_call=1,
    eval_episodes=10,
    debug_summaries=False,
    summarize_grads_and_vars=False,
    seed=None):
  """Trains and evaluates CQL-SAC."""
  logging.info('Training CQL-SAC on: %s', env_name)
  tf.random.set_seed(seed)
  np.random.seed(seed)

  # Load environment.
  env = load_d4rl(env_name)
  tf_env = tf_py_environment.TFPyEnvironment(env)
  strategy = strategy_utils.get_strategy(tpu, use_gpu)

  if not dataset_path.endswith('.tfrecord'):
    dataset_path = os.path.join(dataset_path, env_name,
                                '%s*.tfrecord' % env_name)
  logging.info('Loading dataset from %s', dataset_path)
  dataset_paths = tf.io.gfile.glob(dataset_path)

  # Create dataset.
  with strategy.scope():
    dataset = create_tf_record_dataset(
        dataset_paths,
        batch_size,
        shuffle_buffer_size_per_record=data_shuffle_buffer_size_per_record,
        shuffle_buffer_size=data_shuffle_buffer_size,
        num_shards=data_num_shards,
        cycle_length=data_cycle_length,
        block_length=data_block_length,
        num_parallel_reads=data_parallel_reads,
        num_parallel_calls=data_parallel_calls,
        num_prefetch=data_prefetch,
        strategy=strategy,
        reward_shift=reward_shift,
        action_clipping=action_clipping,
        use_trajectories=use_trajectories)

  # Create agent.
  time_step_spec = tf_env.time_step_spec()
  observation_spec = time_step_spec.observation
  action_spec = tf_env.action_spec()
  with strategy.scope():
    train_step = train_utils.create_train_step()

    actor_net = actor_distribution_network.ActorDistributionNetwork(
        observation_spec,
        action_spec,
        fc_layer_params=actor_fc_layers,
        continuous_projection_net=tanh_normal_projection_network
        .TanhNormalProjectionNetwork)

    critic_net = critic_network.CriticNetwork(
        (observation_spec, action_spec),
        joint_fc_layer_params=critic_joint_fc_layers,
        kernel_initializer='glorot_uniform',
        last_kernel_initializer='glorot_uniform')

    agent = cql_sac_agent.CqlSacAgent(
        time_step_spec,
        action_spec,
        actor_network=actor_net,
        critic_network=critic_net,
        actor_optimizer=tf.keras.optimizers.Adam(
            learning_rate=actor_learning_rate),
        critic_optimizer=tf.keras.optimizers.Adam(
            learning_rate=critic_learning_rate),
        alpha_optimizer=tf.keras.optimizers.Adam(
            learning_rate=alpha_learning_rate),
        cql_alpha=cql_alpha,
        num_cql_samples=num_cql_samples,
        include_critic_entropy_term=include_critic_entropy_term,
        use_lagrange_cql_alpha=use_lagrange_cql_alpha,
        cql_alpha_learning_rate=cql_alpha_learning_rate,
        target_update_tau=5e-3,
        target_update_period=1,
        random_seed=seed,
        cql_tau=cql_tau,
        reward_noise_variance=reward_noise_variance,
        num_bc_steps=bc_steps,
        td_errors_loss_fn=tf.math.squared_difference,
        gamma=0.99,
        reward_scale_factor=reward_scale_factor,
        gradient_clipping=None,
        log_cql_alpha_clipping=log_cql_alpha_clipping,
        softmax_temperature=softmax_temperature,
        debug_summaries=debug_summaries,
        summarize_grads_and_vars=summarize_grads_and_vars,
        train_step_counter=train_step)
    agent.initialize()

  # Create learner.
  saved_model_dir = os.path.join(root_dir, learner.POLICY_SAVED_MODEL_DIR)
  collect_env_step_metric = py_metrics.EnvironmentSteps()
  learning_triggers = [
      triggers.PolicySavedModelTrigger(
          saved_model_dir,
          agent,
          train_step,
          interval=policy_save_interval,
          metadata_metrics={
              triggers.ENV_STEP_METADATA_KEY: collect_env_step_metric
          }),
      triggers.StepPerSecondLogTrigger(train_step, interval=100)
  ]
  cql_learner = learner.Learner(
      root_dir,
      train_step,
      agent,
      experience_dataset_fn=lambda: dataset,
      triggers=learning_triggers,
      summary_interval=summary_interval,
      strategy=strategy)

  # Create actor for evaluation.
  tf_greedy_policy = greedy_policy.GreedyPolicy(agent.policy)
  eval_greedy_policy = py_tf_eager_policy.PyTFEagerPolicy(
      tf_greedy_policy, use_tf_function=True)
  eval_actor = actor.Actor(
      env,
      eval_greedy_policy,
      train_step,
      metrics=actor.eval_metrics(eval_episodes),
      summary_dir=os.path.join(root_dir, 'eval'),
      episodes_per_run=eval_episodes)

  # Run.
  dummy_trajectory = trajectory.mid((), (), (), 0., 1.)
  num_learner_iterations = int(num_gradient_updates /
                               learner_iterations_per_call)
  for _ in range(num_learner_iterations):
    # Mimic collecting environment steps since we loaded a static dataset.
    for _ in range(learner_iterations_per_call):
      collect_env_step_metric(dummy_trajectory)

    cql_learner.run(iterations=learner_iterations_per_call)
    if eval_interval and train_step.numpy() % eval_interval == 0:
      eval_actor.run_and_log()

def main(_):
  logging.set_verbosity(logging.INFO)

  gin.parse_config_files_and_bindings(FLAGS.gin_file, FLAGS.gin_param)

  train_eval(
      root_dir=FLAGS.root_dir,
      dataset_path=FLAGS.dataset_path,
      env_name=FLAGS.env_name,
      tpu=FLAGS.tpu,
      use_gpu=FLAGS.use_gpu,
      num_gradient_updates=FLAGS.num_gradient_updates,
      policy_save_interval=FLAGS.policy_save_interval,
      eval_interval=FLAGS.eval_interval,
      summary_interval=FLAGS.summary_interval,
      learner_iterations_per_call=FLAGS.learner_iterations_per_call,
      use_trajectories=FLAGS.use_trajectories)

if __name__ == '__main__':
  flags.mark_flag_as_required('root_dir')
  flags.mark_flag_as_required('dataset_path')
  app.run(main)
<|endoftext|>


<commit_before># Copyright 2014 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Background tasks for the www_server module.

This module has the logic for handling background tasks for the www frontend.
Long term actions (like periodic tracing), cannot be served synchronously in the
context of a /ajax/endpoint request (would timeout HTTP). Instead, for such long
operations, an instance of |BackgroundTask| is created here and the server
returns just its id. The client can later poll the status of the asynchronous
task to check for its progress.

From a technical viewpoint, each background task is just a python thread
which communicates its progress updates through a Queue. The messages enqueued
are tuples with the following format: (completion_ratio%, 'message string').
"""

import datetime
import Queue
import threading
import time

from memory_inspector.core import backends
from memory_inspector.data import file_storage

_tasks = {}  #id (int) -> |BackgroundTask| instance.

def StartTracer(process, storage_path, interval, count, trace_native_heap):
  assert(isinstance(process, backends.Process))
  task = BackgroundTask(
      TracerMain_,
      storage_path=storage_path,
      backend_name=process.device.backend.name,
      device_id=process.device.id,
      pid=process.pid,
      interval=interval,
      count=count,
      trace_native_heap=trace_native_heap)
  task.start()
  _tasks[task.ident] = task
  return task.ident

def Get(task_id):
  return _tasks.get(task_id)

def TerminateAll():
  for task in _tasks.itervalues():
    if task.is_alive():
      task.terminate()
  _tasks.clear()

def TracerMain_(log, storage_path, backend_name, device_id, pid, interval,
    count, trace_native_heap):
  """Entry point for the background periodic tracer task."""
  # Initialize storage.
  storage = file_storage.Storage(storage_path)

  # Initialize the backend.
  backend = backends.GetBackend(backend_name)
  for k, v in storage.LoadSettings(backend_name).iteritems():
    backend.settings[k] = v

  # Initialize the device.
  device = backends.GetDevice(backend_name, device_id)
  for k, v in storage.LoadSettings(device_id).iteritems():
    device.settings[k] = v

  # Start periodic tracing.
  process = device.GetProcess(pid)
  log.put((1, 'Starting trace (%d dumps x %s s.). Device: %s, process: %s' % (
      count, interval, device.name, process.name)))
  datetime_str = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M')
  archive_name = '%s - %s - %s' % (datetime_str, device.name, process.name)
  archive = storage.OpenArchive(archive_name, create=True)
  heaps_to_symbolize = []

  for i in xrange(1, count + 1):  # [1, count] range is easier to handle.
    process = device.GetProcess(pid)
    if not process:
      log.put((100, 'Process %d died.' % pid))
      return 1
    # Calculate the completion rate proportionally to 80%. We keep the remaining
    # 20% for the final symbolization step (just an approximate estimation).
    completion = 80 * i / count
    log.put((completion, 'Dumping trace %d of %d' % (i, count)))
    archive.StartNewSnapshot()
    # Freeze the process, so that the mmaps and the heap dump are consistent.
    process.Freeze()
    try:
      if trace_native_heap:
        nheap = process.DumpNativeHeap()
        log.put((completion,
                 'Dumped %d native allocations' % len(nheap.allocations)))

      # TODO(primiano): memdump has the bad habit of sending SIGCONT to the
      # process. Fix that, so we are the only one in charge of controlling it.
      mmaps = process.DumpMemoryMaps()
      log.put((completion, 'Dumped %d memory maps' % len(mmaps)))
      archive.StoreMemMaps(mmaps)

      if trace_native_heap:
        nheap.RelativizeStackFrames(mmaps)
        nheap.CalculateResidentSize(mmaps)
        archive.StoreNativeHeap(nheap)
        heaps_to_symbolize += [nheap]
    finally:
      process.Unfreeze()

    if i < count:
      time.sleep(interval)

  if heaps_to_symbolize:
    log.put((90, 'Symbolizing'))
    symbols = backend.ExtractSymbols(
        heaps_to_symbolize, device.settings['native_symbol_paths'] or '')
    expected_symbols_count = len(set.union(
        *[set(x.stack_frames.iterkeys()) for x in heaps_to_symbolize]))
    log.put((99, 'Symbolization complete. Got %d symbols (%.1f%%).' % (
        len(symbols), 100.0 * len(symbols) / expected_symbols_count)))
    archive.StoreSymbols(symbols)

  log.put((100, 'Trace complete.'))
  return 0

class BackgroundTask(threading.Thread):
  def __init__(self, entry_point, *args, **kwargs):
    self._log_queue = Queue.Queue()
    self._progress_log = []  # A list of tuples [(50%, 'msg1'), (100%, 'msg2')].
    super(BackgroundTask, self).__init__(
        target=entry_point,
        args=((self._log_queue,) + args),  # Just propagate all args.
        kwargs=kwargs)

  def run(self):
    try:
      super(BackgroundTask, self).run()
    except Exception, e:
      self._log_queue.put((
          100, 'An error occurred (%s). See the log for more details.' % e))
      raise

  def GetProgress(self):
    """ Returns a tuple (completion_rate, message). """
    while True:
      try:
        self._progress_log += [self._log_queue.get(block=False)]
      except Queue.Empty:
        break
    return self._progress_log
<commit_msg>Use our own task IDs for background tasks in Memory Inspector<commit_after># Copyright 2014 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Background tasks for the www_server module.

This module has the logic for handling background tasks for the www frontend.
Long term actions (like periodic tracing), cannot be served synchronously in the
context of a /ajax/endpoint request (would timeout HTTP). Instead, for such long
operations, an instance of |BackgroundTask| is created here and the server
returns just its id. The client can later poll the status of the asynchronous
task to check for its progress.

From a technical viewpoint, each background task is just a python thread
which communicates its progress updates through a Queue. The messages enqueued
are tuples with the following format: (completion_ratio%, 'message string').
"""

import datetime
import itertools
import Queue
import threading
import time

from memory_inspector.core import backends
from memory_inspector.data import file_storage

_task_id_generator = itertools.count(1)
_tasks = {}  #id (int) -> |BackgroundTask| instance.

def StartTracer(process, storage_path, interval, count, trace_native_heap):
  assert(isinstance(process, backends.Process))
  task = BackgroundTask(
      TracerMain_,
      storage_path=storage_path,
      backend_name=process.device.backend.name,
      device_id=process.device.id,
      pid=process.pid,
      interval=interval,
      count=count,
      trace_native_heap=trace_native_heap)
  task.start()
  task_id = _task_id_generator.next()
  _tasks[task_id] = task
  return task_id

def Get(task_id):
  return _tasks.get(task_id)

def TerminateAll():
  for task in _tasks.itervalues():
    if task.is_alive():
      task.terminate()
  _tasks.clear()

def TracerMain_(log, storage_path, backend_name, device_id, pid, interval,
    count, trace_native_heap):
  """Entry point for the background periodic tracer task."""
  # Initialize storage.
  storage = file_storage.Storage(storage_path)

  # Initialize the backend.
  backend = backends.GetBackend(backend_name)
  for k, v in storage.LoadSettings(backend_name).iteritems():
    backend.settings[k] = v

  # Initialize the device.
  device = backends.GetDevice(backend_name, device_id)
  for k, v in storage.LoadSettings(device_id).iteritems():
    device.settings[k] = v

  # Start periodic tracing.
  process = device.GetProcess(pid)
  log.put((1, 'Starting trace (%d dumps x %s s.). Device: %s, process: %s' % (
      count, interval, device.name, process.name)))
  datetime_str = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M')
  archive_name = '%s - %s - %s' % (datetime_str, device.name, process.name)
  archive = storage.OpenArchive(archive_name, create=True)
  heaps_to_symbolize = []

  for i in xrange(1, count + 1):  # [1, count] range is easier to handle.
    process = device.GetProcess(pid)
    if not process:
      log.put((100, 'Process %d died.' % pid))
      return 1
    # Calculate the completion rate proportionally to 80%. We keep the remaining
    # 20% for the final symbolization step (just an approximate estimation).
    completion = 80 * i / count
    log.put((completion, 'Dumping trace %d of %d' % (i, count)))
    archive.StartNewSnapshot()
    # Freeze the process, so that the mmaps and the heap dump are consistent.
    process.Freeze()
    try:
      if trace_native_heap:
        nheap = process.DumpNativeHeap()
        log.put((completion,
                 'Dumped %d native allocations' % len(nheap.allocations)))

      # TODO(primiano): memdump has the bad habit of sending SIGCONT to the
      # process. Fix that, so we are the only one in charge of controlling it.
      mmaps = process.DumpMemoryMaps()
      log.put((completion, 'Dumped %d memory maps' % len(mmaps)))
      archive.StoreMemMaps(mmaps)

      if trace_native_heap:
        nheap.RelativizeStackFrames(mmaps)
        nheap.CalculateResidentSize(mmaps)
        archive.StoreNativeHeap(nheap)
        heaps_to_symbolize += [nheap]
    finally:
      process.Unfreeze()

    if i < count:
      time.sleep(interval)

  if heaps_to_symbolize:
    log.put((90, 'Symbolizing'))
    symbols = backend.ExtractSymbols(
        heaps_to_symbolize, device.settings['native_symbol_paths'] or '')
    expected_symbols_count = len(set.union(
        *[set(x.stack_frames.iterkeys()) for x in heaps_to_symbolize]))
    log.put((99, 'Symbolization complete. Got %d symbols (%.1f%%).' % (
        len(symbols), 100.0 * len(symbols) / expected_symbols_count)))
    archive.StoreSymbols(symbols)

  log.put((100, 'Trace complete.'))
  return 0

class BackgroundTask(threading.Thread):
  def __init__(self, entry_point, *args, **kwargs):
    self._log_queue = Queue.Queue()
    self._progress_log = []  # A list of tuples [(50%, 'msg1'), (100%, 'msg2')].
    super(BackgroundTask, self).__init__(
        target=entry_point,
        args=((self._log_queue,) + args),  # Just propagate all args.
        kwargs=kwargs)

  def run(self):
    try:
      super(BackgroundTask, self).run()
    except Exception, e:
      self._log_queue.put((
          100, 'An error occurred (%s). See the log for more details.' % e))
      raise

  def GetProgress(self):
    """ Returns a tuple (completion_rate, message). """
    while True:
      try:
        self._progress_log += [self._log_queue.get(block=False)]
      except Queue.Empty:
        break
    return self._progress_log
<|endoftext|>