<commit_before># Copyright 2020 - 2021 MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
A collection of generic interfaces for MONAI transforms.
"""

import warnings
from typing import Any, Callable, Mapping, Optional, Sequence, Union

import numpy as np

from monai.transforms.inverse import InvertibleTransform

# For backwards compatibility (so this still works: from monai.transforms.compose import MapTransform)
from monai.transforms.transform import (  # noqa: F401
    MapTransform,
    Randomizable,
    RandomizableTransform,
    Transform,
    apply_transform,
)
from monai.utils import MAX_SEED, ensure_tuple, get_seed
from monai.utils.enums import TraceKeys

__all__ = ["Compose", "OneOf"]

class Compose(Randomizable, InvertibleTransform):
    """
    ``Compose`` provides the ability to chain a series of callables together in
    a sequential manner. Each transform in the sequence must take a single
    argument and return a single value.

    ``Compose`` can be used in two ways:

    #. With a series of transforms that accept and return a single
       ndarray / tensor / tensor-like parameter.
    #. With a series of transforms that accept and return a dictionary that
       contains one or more parameters. Such transforms must have pass-through
       semantics that unused values in the dictionary must be copied to the return
       dictionary. It is required that the dictionary is copied between input
       and output of each transform.

    If some transform takes a data item dictionary as input, and returns a
    sequence of data items in the transform chain, all following transforms
    will be applied to each item of this list if `map_items` is `True` (the
    default).  If `map_items` is `False`, the returned sequence is passed whole
    to the next callable in the chain.

    For example:

    A `Compose([transformA, transformB, transformC],
    map_items=True)(data_dict)` could achieve the following patch-based
    transformation on the `data_dict` input:

    #. transformA normalizes the intensity of 'img' field in the `data_dict`.
    #. transformB crops out image patches from the 'img' and 'seg' of
       `data_dict`, and return a list of three patch samples::

        {'img': 3x100x100 data, 'seg': 1x100x100 data, 'shape': (100, 100)}
                             applying transformB
                                 ---------->
        [{'img': 3x20x20 data, 'seg': 1x20x20 data, 'shape': (20, 20)},
         {'img': 3x20x20 data, 'seg': 1x20x20 data, 'shape': (20, 20)},
         {'img': 3x20x20 data, 'seg': 1x20x20 data, 'shape': (20, 20)},]

    #. transformC then randomly rotates or flips 'img' and 'seg' of
       each dictionary item in the list returned by transformB.

    The composed transforms will be set the same global random seed if user called
    `set_determinism()`.

    When using the pass-through dictionary operation, you can make use of
    :class:`monai.transforms.adaptors.adaptor` to wrap transforms that don't conform
    to the requirements. This approach allows you to use transforms from
    otherwise incompatible libraries with minimal additional work.

    Note:

        In many cases, Compose is not the best way to create pre-processing
        pipelines. Pre-processing is often not a strictly sequential series of
        operations, and much of the complexity arises when a not-sequential
        set of functions must be called as if it were a sequence.

        Example: images and labels
        Images typically require some kind of normalization that labels do not.
        Both are then typically augmented through the use of random rotations,
        flips, and deformations.
        Compose can be used with a series of transforms that take a dictionary
        that contains 'image' and 'label' entries. This might require wrapping
        `torchvision` transforms before passing them to compose.
        Alternatively, one can create a class with a `__call__` function that
        calls your pre-processing functions taking into account that not all of
        them are called on the labels.
    """

    def __init__(
        self,
        transforms: Optional[Union[Sequence[Callable], Callable]] = None,
        map_items: bool = True,
        unpack_items: bool = False,
    ) -> None:
        if transforms is None:
            transforms = []
        self.transforms = ensure_tuple(transforms)
        self.map_items = map_items
        self.unpack_items = unpack_items
        self.set_random_state(seed=get_seed())

    def set_random_state(self, seed: Optional[int] = None, state: Optional[np.random.RandomState] = None) -> "Compose":
        super().set_random_state(seed=seed, state=state)
        for _transform in self.transforms:
            if not isinstance(_transform, Randomizable):
                continue
            _transform.set_random_state(seed=self.R.randint(MAX_SEED, dtype="uint32"))
        return self

    def randomize(self, data: Optional[Any] = None) -> None:
        for _transform in self.transforms:
            if not isinstance(_transform, Randomizable):
                continue
            try:
                _transform.randomize(data)
            except TypeError as type_error:
                tfm_name: str = type(_transform).__name__
                warnings.warn(
                    f'Transform "{tfm_name}" in Compose not randomized\n{tfm_name}.{type_error}.', RuntimeWarning
                )

    def flatten(self):
        """Return a Composition with a simple list of transforms, as opposed to any nested Compositions.

        e.g., `t1 = Compose([x, x, x, x, Compose([Compose([x, x]), x, x])]).flatten()`
        will result in the equivalent of `t1 = Compose([x, x, x, x, x, x, x, x])`.

        """
        new_transforms = []
        for t in self.transforms:
            if isinstance(t, Compose) and not isinstance(t, OneOf):
                new_transforms += t.flatten().transforms
            else:
                new_transforms.append(t)

        return Compose(new_transforms)

    def __len__(self):
        """Return number of transformations."""
        return len(self.flatten().transforms)

    def __call__(self, input_):
        for _transform in self.transforms:
            input_ = apply_transform(_transform, input_, self.map_items, self.unpack_items)
        return input_

    def inverse(self, data):
        invertible_transforms = [t for t in self.flatten().transforms if isinstance(t, InvertibleTransform)]
        if not invertible_transforms:
            warnings.warn("inverse has been called but no invertible transforms have been supplied")

        # loop backwards over transforms
        for t in reversed(invertible_transforms):
            data = apply_transform(t.inverse, data, self.map_items, self.unpack_items)
        return data

class OneOf(Compose):
    """
    ``OneOf`` provides the ability to randomly choose one transform out of a
    list of callables with pre-defined probabilities for each.

    Args:
        transforms: sequence of callables.
        weights: probabilities corresponding to each callable in transforms.
            Probabilities are normalized to sum to one.

    ``OneOf`` inherits from ``Compose`` and uses args ``map_items`` and ``unpack_items`` in
    the same way.
    """

    def __init__(
        self,
        transforms: Optional[Union[Sequence[Callable], Callable]] = None,
        weights: Optional[Union[Sequence[float], float]] = None,
        map_items: bool = True,
        unpack_items: bool = False,
    ) -> None:
        super().__init__(transforms, map_items, unpack_items)
        if len(self.transforms) == 0:
            weights = []
        elif weights is None or isinstance(weights, float):
            weights = [1.0 / len(self.transforms)] * len(self.transforms)
        if len(weights) != len(self.transforms):
            raise AssertionError("transforms and weights should be same size if both specified as sequences.")
        self.weights = ensure_tuple(self._normalize_probabilities(weights))

    def _normalize_probabilities(self, weights):
        if len(weights) == 0:
            return weights
        weights = np.array(weights)
        if np.any(weights < 0):
            raise AssertionError("Probabilities must be greater than or equal to zero.")
        if np.all(weights == 0):
            raise AssertionError("At least one probability must be greater than zero.")
        weights = weights / weights.sum()
        return list(weights)

    def flatten(self):
        transforms = []
        weights = []
        for t, w in zip(self.transforms, self.weights):
            # if nested, probability is the current weight multiplied by the nested weights,
            # and so on recursively
            if isinstance(t, OneOf):
                tr = t.flatten()
                for t_, w_ in zip(tr.transforms, tr.weights):
                    transforms.append(t_)
                    weights.append(w_ * w)
            else:
                transforms.append(t)
                weights.append(w)
        return OneOf(transforms, weights, self.map_items, self.unpack_items)

    def __call__(self, data):
        if len(self.transforms) == 0:
            return data
        index = self.R.multinomial(1, self.weights).argmax()
        _transform = self.transforms[index]
        data = apply_transform(_transform, data, self.map_items, self.unpack_items)
        # if the data is a mapping (dictionary), append the OneOf transform to the end
        if isinstance(data, Mapping):
            for key in data.keys():
                if self.trace_key(key) in data:
                    self.push_transform(data, key, extra_info={"index": index})
        return data

    def inverse(self, data):
        if len(self.transforms) == 0:
            return data
        if not isinstance(data, Mapping):
            raise RuntimeError("Inverse only implemented for Mapping (dictionary) data")

        # loop until we get an index and then break (since they'll all be the same)
        index = None
        for key in data.keys():
            if self.trace_key(key) in data:
                # get the index of the applied OneOf transform
                index = self.get_most_recent_transform(data, key)[TraceKeys.EXTRA_INFO]["index"]
                # and then remove the OneOf transform
                self.pop_transform(data, key)
        if index is None:
            # no invertible transforms have been applied
            return data

        _transform = self.transforms[index]
        # apply the inverse
        return _transform.inverse(data) if isinstance(_transform, InvertibleTransform) else data
<commit_msg>[DLMED] add missing doc-string (#3646)<commit_after># Copyright 2020 - 2021 MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
A collection of generic interfaces for MONAI transforms.
"""

import warnings
from typing import Any, Callable, Mapping, Optional, Sequence, Union

import numpy as np

from monai.transforms.inverse import InvertibleTransform

# For backwards compatibility (so this still works: from monai.transforms.compose import MapTransform)
from monai.transforms.transform import (  # noqa: F401
    MapTransform,
    Randomizable,
    RandomizableTransform,
    Transform,
    apply_transform,
)
from monai.utils import MAX_SEED, ensure_tuple, get_seed
from monai.utils.enums import TraceKeys

__all__ = ["Compose", "OneOf"]

class Compose(Randomizable, InvertibleTransform):
    """
    ``Compose`` provides the ability to chain a series of callables together in
    a sequential manner. Each transform in the sequence must take a single
    argument and return a single value.

    ``Compose`` can be used in two ways:

    #. With a series of transforms that accept and return a single
       ndarray / tensor / tensor-like parameter.
    #. With a series of transforms that accept and return a dictionary that
       contains one or more parameters. Such transforms must have pass-through
       semantics that unused values in the dictionary must be copied to the return
       dictionary. It is required that the dictionary is copied between input
       and output of each transform.

    If some transform takes a data item dictionary as input, and returns a
    sequence of data items in the transform chain, all following transforms
    will be applied to each item of this list if `map_items` is `True` (the
    default).  If `map_items` is `False`, the returned sequence is passed whole
    to the next callable in the chain.

    For example:

    A `Compose([transformA, transformB, transformC],
    map_items=True)(data_dict)` could achieve the following patch-based
    transformation on the `data_dict` input:

    #. transformA normalizes the intensity of 'img' field in the `data_dict`.
    #. transformB crops out image patches from the 'img' and 'seg' of
       `data_dict`, and return a list of three patch samples::

        {'img': 3x100x100 data, 'seg': 1x100x100 data, 'shape': (100, 100)}
                             applying transformB
                                 ---------->
        [{'img': 3x20x20 data, 'seg': 1x20x20 data, 'shape': (20, 20)},
         {'img': 3x20x20 data, 'seg': 1x20x20 data, 'shape': (20, 20)},
         {'img': 3x20x20 data, 'seg': 1x20x20 data, 'shape': (20, 20)},]

    #. transformC then randomly rotates or flips 'img' and 'seg' of
       each dictionary item in the list returned by transformB.

    The composed transforms will be set the same global random seed if user called
    `set_determinism()`.

    When using the pass-through dictionary operation, you can make use of
    :class:`monai.transforms.adaptors.adaptor` to wrap transforms that don't conform
    to the requirements. This approach allows you to use transforms from
    otherwise incompatible libraries with minimal additional work.

    Note:

        In many cases, Compose is not the best way to create pre-processing
        pipelines. Pre-processing is often not a strictly sequential series of
        operations, and much of the complexity arises when a not-sequential
        set of functions must be called as if it were a sequence.

        Example: images and labels
        Images typically require some kind of normalization that labels do not.
        Both are then typically augmented through the use of random rotations,
        flips, and deformations.
        Compose can be used with a series of transforms that take a dictionary
        that contains 'image' and 'label' entries. This might require wrapping
        `torchvision` transforms before passing them to compose.
        Alternatively, one can create a class with a `__call__` function that
        calls your pre-processing functions taking into account that not all of
        them are called on the labels.

    Args:
        transforms: sequence of callables.
        map_items: whether to apply transform to each item in the input `data` if `data` is a list or tuple.
            defaults to `True`.
        unpack_items: whether to unpack input `data` with `*` as parameters for the callable function of transform.
            defaults to `False`.

    """

    def __init__(
        self,
        transforms: Optional[Union[Sequence[Callable], Callable]] = None,
        map_items: bool = True,
        unpack_items: bool = False,
    ) -> None:
        if transforms is None:
            transforms = []
        self.transforms = ensure_tuple(transforms)
        self.map_items = map_items
        self.unpack_items = unpack_items
        self.set_random_state(seed=get_seed())

    def set_random_state(self, seed: Optional[int] = None, state: Optional[np.random.RandomState] = None) -> "Compose":
        super().set_random_state(seed=seed, state=state)
        for _transform in self.transforms:
            if not isinstance(_transform, Randomizable):
                continue
            _transform.set_random_state(seed=self.R.randint(MAX_SEED, dtype="uint32"))
        return self

    def randomize(self, data: Optional[Any] = None) -> None:
        for _transform in self.transforms:
            if not isinstance(_transform, Randomizable):
                continue
            try:
                _transform.randomize(data)
            except TypeError as type_error:
                tfm_name: str = type(_transform).__name__
                warnings.warn(
                    f'Transform "{tfm_name}" in Compose not randomized\n{tfm_name}.{type_error}.', RuntimeWarning
                )

    def flatten(self):
        """Return a Composition with a simple list of transforms, as opposed to any nested Compositions.

        e.g., `t1 = Compose([x, x, x, x, Compose([Compose([x, x]), x, x])]).flatten()`
        will result in the equivalent of `t1 = Compose([x, x, x, x, x, x, x, x])`.

        """
        new_transforms = []
        for t in self.transforms:
            if isinstance(t, Compose) and not isinstance(t, OneOf):
                new_transforms += t.flatten().transforms
            else:
                new_transforms.append(t)

        return Compose(new_transforms)

    def __len__(self):
        """Return number of transformations."""
        return len(self.flatten().transforms)

    def __call__(self, input_):
        for _transform in self.transforms:
            input_ = apply_transform(_transform, input_, self.map_items, self.unpack_items)
        return input_

    def inverse(self, data):
        invertible_transforms = [t for t in self.flatten().transforms if isinstance(t, InvertibleTransform)]
        if not invertible_transforms:
            warnings.warn("inverse has been called but no invertible transforms have been supplied")

        # loop backwards over transforms
        for t in reversed(invertible_transforms):
            data = apply_transform(t.inverse, data, self.map_items, self.unpack_items)
        return data

class OneOf(Compose):
    """
    ``OneOf`` provides the ability to randomly choose one transform out of a
    list of callables with pre-defined probabilities for each.

    Args:
        transforms: sequence of callables.
        weights: probabilities corresponding to each callable in transforms.
            Probabilities are normalized to sum to one.
        map_items: whether to apply transform to each item in the input `data` if `data` is a list or tuple.
            defaults to `True`.
        unpack_items: whether to unpack input `data` with `*` as parameters for the callable function of transform.
            defaults to `False`.

    """

    def __init__(
        self,
        transforms: Optional[Union[Sequence[Callable], Callable]] = None,
        weights: Optional[Union[Sequence[float], float]] = None,
        map_items: bool = True,
        unpack_items: bool = False,
    ) -> None:
        super().__init__(transforms, map_items, unpack_items)
        if len(self.transforms) == 0:
            weights = []
        elif weights is None or isinstance(weights, float):
            weights = [1.0 / len(self.transforms)] * len(self.transforms)
        if len(weights) != len(self.transforms):
            raise AssertionError("transforms and weights should be same size if both specified as sequences.")
        self.weights = ensure_tuple(self._normalize_probabilities(weights))

    def _normalize_probabilities(self, weights):
        if len(weights) == 0:
            return weights
        weights = np.array(weights)
        if np.any(weights < 0):
            raise AssertionError("Probabilities must be greater than or equal to zero.")
        if np.all(weights == 0):
            raise AssertionError("At least one probability must be greater than zero.")
        weights = weights / weights.sum()
        return list(weights)

    def flatten(self):
        transforms = []
        weights = []
        for t, w in zip(self.transforms, self.weights):
            # if nested, probability is the current weight multiplied by the nested weights,
            # and so on recursively
            if isinstance(t, OneOf):
                tr = t.flatten()
                for t_, w_ in zip(tr.transforms, tr.weights):
                    transforms.append(t_)
                    weights.append(w_ * w)
            else:
                transforms.append(t)
                weights.append(w)
        return OneOf(transforms, weights, self.map_items, self.unpack_items)

    def __call__(self, data):
        if len(self.transforms) == 0:
            return data
        index = self.R.multinomial(1, self.weights).argmax()
        _transform = self.transforms[index]
        data = apply_transform(_transform, data, self.map_items, self.unpack_items)
        # if the data is a mapping (dictionary), append the OneOf transform to the end
        if isinstance(data, Mapping):
            for key in data.keys():
                if self.trace_key(key) in data:
                    self.push_transform(data, key, extra_info={"index": index})
        return data

    def inverse(self, data):
        if len(self.transforms) == 0:
            return data
        if not isinstance(data, Mapping):
            raise RuntimeError("Inverse only implemented for Mapping (dictionary) data")

        # loop until we get an index and then break (since they'll all be the same)
        index = None
        for key in data.keys():
            if self.trace_key(key) in data:
                # get the index of the applied OneOf transform
                index = self.get_most_recent_transform(data, key)[TraceKeys.EXTRA_INFO]["index"]
                # and then remove the OneOf transform
                self.pop_transform(data, key)
        if index is None:
            # no invertible transforms have been applied
            return data

        _transform = self.transforms[index]
        # apply the inverse
        return _transform.inverse(data) if isinstance(_transform, InvertibleTransform) else data
<|endoftext|>


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

### Implementation of the DDPG algorithm [1].
###
### [1] https://arxiv.org/pdf/1509.02971.pdf

from ai_utils import noise, replay_buffer as rb
from ai_utils.actor import actor_network
from ai_utils.critic import critic_network

from keras import backend as K
import tensorflow as tf
import numpy as np

## TODO: -> receive simulator to apply operations at environment
##
##       -> integrate with robot environment

class ddpg:
    ##
    ## definition of our DDPG algorithm, default values are according to our
    ## reference [1].
    ##
    ##  -> action_dim:  total of actions;
    ##  -> state_dim:   total of states;
    ##  -> batch size:  batch size for sampling;
    ##
    ##  hyper parameters:
    ##  -> wd:      weight decay for Q;
    ##  -> gamma:   discount factor;
    ##  -> tau:     soft target updates
    ##  -> lra:     learning rate for ACTOR
    ##  -> lrc:     learning rate for CRITIC
    ##
    def __init__(self, robot, state_dim, action_dim,
        batch_size=32, wd=.01, gamma=.99, tau=.001, lra=.0001, lrc=.001,
        path='./aidata/'):

        self.robot = robot
        self.path = path
        self.noise = noise.ou()

        self.batch_size = batch_size

        self.wd = wd
        self.gamma = gamma
        self.tau = tau

        self.lra = lra
        self.lrc = lrc

        self.action_dim = action_dim
        self.state_dim = state_dim

        # gpu usage
        # config = tf.ConfigProto()
        # config.gpu_options.allow_growth = True

        self.sess = tf.Session() # config=config
        K.set_session(self.sess)

    def load_weights(self, actor, critic):
        actor.load_weights(self.path)
        critic.load_weights(self.path)

    ##########
    ## act!
    ##########

    ## train our ddpg model
    def train(self, max_episodes, max_steps, buffer_size=1000000):
        # initialize actor and critic networks
        actor = actor_network(self.sess, self.state_dim, self.action_dim,
            self.batch_size, self.tau, self.lra)

        critic = critic_network(self.sess, self.state_dim, self.action_dim,
            self.batch_size, self.tau, self.lrc)

        # load weights
        self.load_weights(actor, critic)

        # initialize replay buffer R
        buff = rb.replay_buffer(buffer_size, self.batch_size)

        # helpers at training
        epsilon = 1
        explore_decay = 1.0/100000.

        for episode in range(max_episodes):
            ## initialize a random process for action exploration from our
            ## VREP environment
            # self.simulator.restart()

            ## get initial state from our VREP environment
            s_t = self.bake(self.robot.get_state())

            total_reward = 0.
            for t in range(max_steps):
                ## select action according to our current policy
                a_t = np.zeros(self.action_dim)
                noise_t = np.zeros(self.action_dim)

                a_t_raw = actor.helper.predict(s_t)

                # apply exploration noise
                epsilon = max(epsilon-explore_decay, 0)
                noise_t = epsilon*self.noise.apply(a_t_raw)

                a_t = a_t_raw + noise_t

                ## execute action and observe our new reward+state
                new_s_t, r_t, done = self.robot.act(a_t)
                new_s_t = self.bake(new_s_t)

                ## store transition at our buffer
                buff.store(s_t, a_t, r_t, new_s_t, done)

                ## sample random minibatch of transitions from buffer
                states, actions, rewards, new_states, dones = buff.get_batch()

                ## set y according to our choice
                target_q_values = critic.target.predict([new_states, \
                    actor.target.predict(new_states)])

                y_t = r_t + dones*self.gamma*target_q_values

                ## update critics by minimizing loss
                loss = critic.helper.train_on_batch([states, actions], y_t)

                ## update actor policy using the sampled policy gradient
                gradient_actions = actor.helper.predict(states)
                gradients = critic.gradients(states, gradient_actions)

                actor.train(states, gradients)

                ## update target networks
                actor.target_train()
                critic.target_train()

                s_t = new_s_t
                total_reward += r_t

                ## did we end our episode?
                if done:
                    break

            if episode % 5 is 0:
                actor.save_weights(self.path)
                critic.save_weights(self.path)

            print '*********************************************'
            print 'EPISODE: ' + str(episode)
            print '\tTOTAL REWARD: ' + str(total_reward)
            print '*********************************************'

    ## run our ddpg model
    def tick(self, max_episodes, max_steps):
        # initialize actor and critic networks
        actor = actor_network(self.sess, self.state_dim, self.action_dim)
        critic = critic_network(self.sess, self.state_dim, self.action_dim)

        # load weights
        self.load_weights(actor, critic)

        for episode in range(max_episodes):
            ## initialize a random process for action exploration from our
            ## VREP environment
            # self.simulator.restart()

            ## get initial state from our VREP environment
            s_t = self.bake(self.robot.get_state())

            total_reward = 0.
            for t in range(max_steps):
                ## select action according to our current policy
                a_t = actor.helper.predict(s_t)

                ## execute action and observe our new reward+state
                new_s_t, r_t, done = self.robot.act(a_t)

                s_t = new_s_t
                total_reward += r_t

                ## did we end our episode?
                if done:
                    break

            print '*********************************************'
            print '\tEPISODE: ' + str(episode)
            print '\tTOTAL REWARD: ' + str(total_reward)
            print '*********************************************'

    ##########
    ## helpers
    ##########

    ## bake our input
    def bake(self, arr):
        return np.reshape(arr, (1, len(arr))).astype(np.float32)

## testing...
class robot:
    def __init__(self, s_dim, a_dim):
        self.s = s_dim
        self.a = a_dim

    def get_state(self):
        return [1]*25

    def act(self, actions):
        print '\t\t-> Actions to be executed: ' + str(actions)
        return [1]*self.s, 1, False

if __name__ == "__main__":
    r = robot(25, 8)
    ai = ddpg(r, 25, 8)
    ai.train(10, 10)
<commit_msg>changed method name.<commit_after># -*- coding: utf-8 -*-

### Implementation of the DDPG algorithm [1].
###
### [1] https://arxiv.org/pdf/1509.02971.pdf

from ai_utils import noise, replay_buffer as rb
from ai_utils.actor import actor_network
from ai_utils.critic import critic_network

from keras import backend as K
import tensorflow as tf
import numpy as np

## TODO: -> receive simulator to apply operations at environment
##
##       -> integrate with robot environment

class ddpg:
    ##
    ## definition of our DDPG algorithm, default values are according to our
    ## reference [1].
    ##
    ##  -> action_dim:  total of actions;
    ##  -> state_dim:   total of states;
    ##  -> batch size:  batch size for sampling;
    ##
    ##  hyper parameters:
    ##  -> wd:      weight decay for Q;
    ##  -> gamma:   discount factor;
    ##  -> tau:     soft target updates
    ##  -> lra:     learning rate for ACTOR
    ##  -> lrc:     learning rate for CRITIC
    ##
    def __init__(self, robot, state_dim, action_dim,
        batch_size=32, wd=.01, gamma=.99, tau=.001, lra=.0001, lrc=.001,
        path='./aidata/'):

        self.robot = robot
        self.path = path
        self.noise = noise.ou()

        self.batch_size = batch_size

        self.wd = wd
        self.gamma = gamma
        self.tau = tau

        self.lra = lra
        self.lrc = lrc

        self.action_dim = action_dim
        self.state_dim = state_dim

        # gpu usage
        # config = tf.ConfigProto()
        # config.gpu_options.allow_growth = True

        self.sess = tf.Session() # config=config
        K.set_session(self.sess)

    def load_weights(self, actor, critic):
        actor.load_weights(self.path)
        critic.load_weights(self.path)

    ##########
    ## act!
    ##########

    ## train our ddpg model
    def train(self, max_episodes, max_steps, buffer_size=1000000):
        # initialize actor and critic networks
        actor = actor_network(self.sess, self.state_dim, self.action_dim,
            self.batch_size, self.tau, self.lra)

        critic = critic_network(self.sess, self.state_dim, self.action_dim,
            self.batch_size, self.tau, self.lrc)

        # load weights
        self.load_weights(actor, critic)

        # initialize replay buffer R
        buff = rb.replay_buffer(buffer_size, self.batch_size)

        # helpers at training
        epsilon = 1
        explore_decay = 1.0/100000.

        for episode in range(max_episodes):
            ## initialize a random process for action exploration from our
            ## VREP environment
            # self.simulator.restart()

            ## get initial state from our VREP environment
            s_t = self.bake(self.robot.get_state())

            total_reward = 0.
            for t in range(max_steps):
                ## select action according to our current policy
                a_t = np.zeros(self.action_dim)
                noise_t = np.zeros(self.action_dim)

                a_t_raw = actor.helper.predict(s_t)

                # apply exploration noise
                epsilon = max(epsilon-explore_decay, 0)
                noise_t = epsilon*self.noise.apply(a_t_raw)

                a_t = a_t_raw + noise_t

                ## execute action and observe our new reward+state
                new_s_t, r_t, done = self.robot.act(a_t)
                new_s_t = self.bake(new_s_t)

                ## store transition at our buffer
                buff.store(s_t, a_t, r_t, new_s_t, done)

                ## sample random minibatch of transitions from buffer
                states, actions, rewards, new_states, dones = buff.get_batch()

                ## set y according to our choice
                target_q_values = critic.target.predict([new_states, \
                    actor.target.predict(new_states)])

                y_t = r_t + dones*self.gamma*target_q_values

                ## update critics by minimizing loss
                loss = critic.helper.train_on_batch([states, actions], y_t)

                ## update actor policy using the sampled policy gradient
                gradient_actions = actor.helper.predict(states)
                gradients = critic.gradients(states, gradient_actions)

                actor.train(states, gradients)

                ## update target networks
                actor.target_train()
                critic.target_train()

                s_t = new_s_t
                total_reward += r_t

                ## did we end our episode?
                if done:
                    break

            if episode % 5 is 0:
                actor.save_weights(self.path)
                critic.save_weights(self.path)

            print '*********************************************'
            print 'EPISODE: ' + str(episode)
            print '\tTOTAL REWARD: ' + str(total_reward)
            print '*********************************************'

    ## run our ddpg model
    def run(self, max_episodes, max_steps):
        # initialize actor and critic networks
        actor = actor_network(self.sess, self.state_dim, self.action_dim)
        critic = critic_network(self.sess, self.state_dim, self.action_dim)

        # load weights
        self.load_weights(actor, critic)

        for episode in range(max_episodes):
            ## initialize a random process for action exploration from our
            ## VREP environment
            # self.simulator.restart()

            ## get initial state from our VREP environment
            s_t = self.bake(self.robot.get_state())

            total_reward = 0.
            for t in range(max_steps):
                ## select action according to our current policy
                a_t = actor.helper.predict(s_t)

                ## execute action and observe our new reward+state
                new_s_t, r_t, done = self.robot.act(a_t)

                s_t = new_s_t
                total_reward += r_t

                ## did we end our episode?
                if done:
                    break

            print '*********************************************'
            print '\tEPISODE: ' + str(episode)
            print '\tTOTAL REWARD: ' + str(total_reward)
            print '*********************************************'

    ##########
    ## helpers
    ##########

    ## bake our input
    def bake(self, arr):
        return np.reshape(arr, (1, len(arr))).astype(np.float32)

## testing...
class robot:
    def __init__(self, s_dim, a_dim):
        self.s = s_dim
        self.a = a_dim

    def get_state(self):
        return [1]*25

    def act(self, actions):
        print '\t\t-> Actions to be executed: ' + str(actions)
        return [1]*self.s, 1, False

if __name__ == "__main__":
    r = robot(25, 8)
    ai = ddpg(r, 25, 8)
    ai.train(10, 10)
<|endoftext|>


<commit_before>from ._config import PluginConfiguration
from ._parser import EndpointParser
from treq.client import HTTPClient
from twisted.internet import reactor, defer
from twisted.internet.interfaces import IHalfCloseableProtocol
from twisted.python import log
from twisted.python.failure import Failure
from twisted.web import server, proxy
from twisted.web.client import Agent
from twisted.web.resource import Resource
from twisted.web.server import NOT_DONE_YET
from urllib import quote as urlquote
from zope.interface import directlyProvides
import StringIO
import json
import treq
import urlparse

class NoPostHooks(Exception):
    """
    Do not run any post-hooks, because of an incompatible Docker response type
    (streaming/hijacked or chunked).
    """

class DockerProxyClient(proxy.ProxyClient):
    """
    An HTTP proxy which knows how to break HTTP just right so that Docker
    stream (attach/events) API calls work.

    self.http: A boolean which reflects whether the connection is in HTTP mode
        (True) or "hijack" mode (False). See
        https://docs.docker.com/reference/api/docker_remote_api_v1.14/#32-hijacking
    """

    http = True
    _streaming = False
    _listener = None
    _responsePartBuffer = b""

    def _fireListener(self, result):
        if self._listener is not None:
            d = self._listener
            self._listener = None
            d.callback(result)

    def setStreamingMode(self, streamingMode):
        """
        Allow anyone with a reference to us to toggle on/off streaming mode.
        Useful when we have no post-hooks (and no indication from Docker that
        it's sending packets of JSON e.g. with build) and we want to avoid
        buffering slow responses in memory.
        """
        self._streaming = streamingMode
        if streamingMode:
            self._fireListener(Failure(NoPostHooks()))

    def registerListener(self, d):
        """
        Register a one shot listener, which can fire either with:
           * Failure(NoPostHooks()) if the proxy is handling comms back to
             the client (streaming/chunked modes), or
           * A tuple containing the (response, code, content-type).
        """
        self._listener = d

    def _handleRawStream(self):
        """
        Switch the current connection to be a "hijacked" aka raw stream: one
        where bytes just get naively proxied back and forth.
        """
        def loseWriteConnectionReason(reason):
            # discard the reason, for compatibility with readConnectionLost
            self.transport.loseWriteConnection()
        self.father.transport.readConnectionLost = loseWriteConnectionReason
        directlyProvides(self.father.transport, IHalfCloseableProtocol)
        self.http = False
        self.father.transport.write(
            "HTTP/1.1 200 OK\r\n"
            "Content-Type: application/vnd.docker.raw-stream\r\n"
            "\r\n")
        def stdinHandler(data):
            self.transport.write(data)
        self.father.transport.protocol.dataReceived = stdinHandler
        self.setStreamingMode(True)

    def handleHeader(self, key, value):
        # print key, "=>", value
        if (key.lower() == "content-type" and
                value == "application/vnd.docker.raw-stream"):
            self._handleRawStream()
        # XXX Turns out, the build endpoint doesn't actually used chunked
        # encoding. It just sends some JSON documents which maybe happen to
        # line up with packet boundaries. So the following if statement is both
        # untested and potentially never triggered in practice. :(
        if key.lower() == "transfer-encoding" and value == "chunked":
            self.setStreamingMode(True)
        return proxy.ProxyClient.handleHeader(self, key, value)

    def handleResponsePart(self, buffer):
        """
        If we're not in streaming mode, buffer the response part(s).
        """
        if self._streaming:
            proxy.ProxyClient.handleResponsePart(self, buffer)
        else:
            # print "BUFFER!", buffer
            self._responsePartBuffer += buffer

    def handleResponseEnd(self):
        """
        If we're completing a chunked response, up-call to handle it like a
        regular reverse proxy.

        If we're completing a non-chunked response, fire the post-hooks.

        If we're completing a hijacked response, pass through the connection
        close.
        """
        if self.http:
            if self._streaming:
                return proxy.ProxyClient.handleResponseEnd(self)
            else:
                contentType = self.father.responseHeaders.getRawHeaders("content-type")
                if contentType:
                    contentType = contentType[0]
                else:
                    contentType = None
                body = self._responsePartBuffer
                self._fireListener(
                        {"PowerstripProtocolVersion": 1,
                         "ModifiedServerResponse":
                            {"Body": body,
                             "Code": self.father.code,
                             "ContentType": contentType}})
        else:
            self.father.transport.loseConnection()

    def rawDataReceived(self, data):
        if self.http:
            return proxy.ProxyClient.rawDataReceived(self, data)
        self.father.transport.write(data)

class DockerProxyClientFactory(proxy.ProxyClientFactory):
    protocol = DockerProxyClient
    _listener = None
    def onCreate(self, d):
        self._listener = d

    def _fireListener(self, result):
        if self._listener is not None:
            d = self._listener
            self._listener = None
            d.callback(result)

    def buildProtocol(self, addr):
        client = proxy.ProxyClientFactory.buildProtocol(self, addr)
        self._fireListener(client)
        return client

class DockerProxy(proxy.ReverseProxyResource):
    proxyClientFactoryClass = DockerProxyClientFactory

    def __init__(self, dockerAddr=None, dockerPort=None, dockerSocket=None,
            path='', reactor=reactor, config=None):
        """
        A docker proxy resource which knows how to connect to real Docker
        daemon either via socket (dockerSocket specified) or address + port for
        TCP connection (dockerAddr + dockerPort specified).
        """
        if config is None:
            # Try to get the configuration from the default place on the
            # filesystem.
            self.config = PluginConfiguration()
        else:
            self.config = config
        self.config.read_and_parse()
        self.parser = EndpointParser(self.config)
        Resource.__init__(self)
        self.host = dockerAddr
        self.port = dockerPort
        self.socket = dockerSocket
        self.path = path
        self.reactor = reactor
        proxy.ReverseProxyResource.__init__(self, dockerAddr, dockerPort, path, reactor) # NB dockerAddr is not actually used
        self.agent = Agent(reactor) # no connectionpool
        self.client = HTTPClient(self.agent)

    def render(self, request, reactor=reactor):
        # We are processing a leaf request.
        # Get the original request body from the client.
        skipPreHooks = False
        if request.requestHeaders.getRawHeaders('content-type') == ["application/json"]:
            originalRequestBody = request.content.read()
            request.content.seek(0) # hee hee
        elif request.requestHeaders.getRawHeaders('content-type') == ["application/tar"]:
            # XXX We can't JSON encode binary data, so don't even try.
            skipPreHooks = True
            originalRequestBody = None
        else:
            originalRequestBody = None
        preHooks = []
        postHooks = []
        d = defer.succeed(None)
        for endpoint in self.parser.match_endpoint(request.method, request.uri.split("?")[0]):
            # It's possible for a request to match multiple endpoint
            # definitions.  Order of matched endpoint is not defined in
            # that case.
            adapters = self.config.endpoint(endpoint)
            preHooks.extend(adapters.pre)
            postHooks.extend(adapters.post)
        def callPreHook(result, hookURL):
            if result is None:
                newRequestBody = originalRequestBody
            else:
                newRequestBody = result["ModifiedClientRequest"]["Body"]
            return self.client.post(hookURL, json.dumps({
                        "PowerstripProtocolVersion": 1,
                        "Type": "pre-hook",
                        "ClientRequest": {
                            "Method": request.method,
                            "Request": request.uri,
                            "Body": newRequestBody,
                            # XXX This would need a ContentType header... if we
                            # were to support non-JSON POST bodies, like build
                            # contexts in pre-hooks.
                        }
                    }), headers={'Content-Type': ['application/json']})
        if not skipPreHooks:
            for preHook in preHooks:
                hookURL = self.config.adapter_uri(preHook)
                d.addCallback(callPreHook, hookURL=hookURL)
                d.addCallback(treq.json_content)
        def doneAllPrehooks(result):
            # Finally pass through the request to actual Docker.  For now we
            # mutate request in-place in such a way that ReverseProxyResource
            # understands it.
            if result is not None:
                requestBody = b""
                bodyFromAdapter = result["ModifiedClientRequest"]["Body"]
                if bodyFromAdapter is not None:
                    requestBody = bodyFromAdapter.encode("utf-8")
                request.content = StringIO.StringIO(requestBody)
                request.requestHeaders.setRawHeaders(b"content-length",
                        [str(len(requestBody))])
            ###########################
            # The following code is copied from t.w.proxy.ReverseProxy so that
            # clientFactory reference can be kept.
            if not self.socket:
                if self.port == 80:
                    host = self.host
                else:
                    host = "%s:%d" % (self.host, self.port)
                request.requestHeaders.setRawHeaders(b"host", [host])
            request.content.seek(0, 0)
            qs = urlparse.urlparse(request.uri)[4]
            if qs:
                rest = self.path + '?' + qs
            else:
                rest = self.path
            allRequestHeaders = request.getAllHeaders()
            if ("transfer-encoding" in allRequestHeaders
                    and allRequestHeaders["transfer-encoding"] == "chunked"):
                del allRequestHeaders["transfer-encoding"]
            # XXX Streaming the contents of the request body into memory could
            # cause OOM issues for large build contexts POSTed through
            # powerstrip.
            body = request.content.read()
            allRequestHeaders["content-length"] = str(len(body))
            clientFactory = self.proxyClientFactoryClass(
                request.method, rest, request.clientproto,
                allRequestHeaders, body, request)
            ###########################
            if self.socket:
                self.reactor.connectUNIX(self.socket, clientFactory)
            else:
                self.reactor.connectTCP(self.host, self.port, clientFactory)
            d = defer.Deferred()
            clientFactory.onCreate(d)
            return d
        d.addCallback(doneAllPrehooks)
        def inspect(client):
            # If there are no post-hooks, allow the response to be streamed
            # back to the client, rather than buffered.
            d = defer.Deferred()
            client.registerListener(d)
            if not postHooks:
                client.setStreamingMode(True)
            return d
        d.addCallback(inspect)
        def callPostHook(result, hookURL):
            serverResponse = result["ModifiedServerResponse"]
            return self.client.post(hookURL, json.dumps({
                        # TODO Write tests for the information provided to the adapter.
                        "PowerstripProtocolVersion": 1,
                        "Type": "post-hook",
                        "ClientRequest": {
                            "Method": request.method,
                            "Request": request.uri,
                            "Body": originalRequestBody,
                            },
                        "ServerResponse": {
                            "ContentType": serverResponse["ContentType"],
                            "Body": serverResponse["Body"],
                            "Code": serverResponse["Code"],
                        },
                    }), headers={'Content-Type': ['application/json']})
        # XXX Need to skip post-hooks for tar archives from e.g. docker export.
        for postHook in postHooks:
            hookURL = self.config.adapter_uri(postHook)
            d.addCallback(callPostHook, hookURL=hookURL)
            d.addCallback(treq.json_content)
        def sendFinalResponseToClient(result):
            # Write the final response to the client.
            request.write(result["ModifiedServerResponse"]["Body"].encode("utf-8"))
            request.finish()
        d.addCallback(sendFinalResponseToClient)
        def squashNoPostHooks(failure):
            failure.trap(NoPostHooks)
        d.addErrback(squashNoPostHooks)
        d.addErrback(log.err, 'while running chain')
        return NOT_DONE_YET

    def getChild(self, path, request):
        fragments = request.uri.split("/")
        fragments.pop(0)
        proxyArgs = (self.host, self.port, self.socket, self.path + '/' + urlquote(path, safe=""),
                     self.reactor)
        #if not request.postpath:
        resource = DockerProxy(*proxyArgs, config=self.config)
        return resource

class ServerProtocolFactory(server.Site):
    def __init__(self, dockerAddr=None, dockerPort=None, dockerSocket=None, config=None):
        self.root = DockerProxy(dockerAddr, dockerPort, dockerSocket, config=config)
        server.Site.__init__(self, self.root)
<commit_msg>Remove XXX - this isn't an addresable concern.<commit_after>from ._config import PluginConfiguration
from ._parser import EndpointParser
from treq.client import HTTPClient
from twisted.internet import reactor, defer
from twisted.internet.interfaces import IHalfCloseableProtocol
from twisted.python import log
from twisted.python.failure import Failure
from twisted.web import server, proxy
from twisted.web.client import Agent
from twisted.web.resource import Resource
from twisted.web.server import NOT_DONE_YET
from urllib import quote as urlquote
from zope.interface import directlyProvides
import StringIO
import json
import treq
import urlparse

class NoPostHooks(Exception):
    """
    Do not run any post-hooks, because of an incompatible Docker response type
    (streaming/hijacked or chunked).
    """

class DockerProxyClient(proxy.ProxyClient):
    """
    An HTTP proxy which knows how to break HTTP just right so that Docker
    stream (attach/events) API calls work.

    self.http: A boolean which reflects whether the connection is in HTTP mode
        (True) or "hijack" mode (False). See
        https://docs.docker.com/reference/api/docker_remote_api_v1.14/#32-hijacking
    """

    http = True
    _streaming = False
    _listener = None
    _responsePartBuffer = b""

    def _fireListener(self, result):
        if self._listener is not None:
            d = self._listener
            self._listener = None
            d.callback(result)

    def setStreamingMode(self, streamingMode):
        """
        Allow anyone with a reference to us to toggle on/off streaming mode.
        Useful when we have no post-hooks (and no indication from Docker that
        it's sending packets of JSON e.g. with build) and we want to avoid
        buffering slow responses in memory.
        """
        self._streaming = streamingMode
        if streamingMode:
            self._fireListener(Failure(NoPostHooks()))

    def registerListener(self, d):
        """
        Register a one shot listener, which can fire either with:
           * Failure(NoPostHooks()) if the proxy is handling comms back to
             the client (streaming/chunked modes), or
           * A tuple containing the (response, code, content-type).
        """
        self._listener = d

    def _handleRawStream(self):
        """
        Switch the current connection to be a "hijacked" aka raw stream: one
        where bytes just get naively proxied back and forth.
        """
        def loseWriteConnectionReason(reason):
            # discard the reason, for compatibility with readConnectionLost
            self.transport.loseWriteConnection()
        self.father.transport.readConnectionLost = loseWriteConnectionReason
        directlyProvides(self.father.transport, IHalfCloseableProtocol)
        self.http = False
        self.father.transport.write(
            "HTTP/1.1 200 OK\r\n"
            "Content-Type: application/vnd.docker.raw-stream\r\n"
            "\r\n")
        def stdinHandler(data):
            self.transport.write(data)
        self.father.transport.protocol.dataReceived = stdinHandler
        self.setStreamingMode(True)

    def handleHeader(self, key, value):
        # print key, "=>", value
        if (key.lower() == "content-type" and
                value == "application/vnd.docker.raw-stream"):
            self._handleRawStream()
        # XXX Turns out, the build endpoint doesn't actually used chunked
        # encoding. It just sends some JSON documents which maybe happen to
        # line up with packet boundaries. So the following if statement is both
        # untested and potentially never triggered in practice. :(
        if key.lower() == "transfer-encoding" and value == "chunked":
            self.setStreamingMode(True)
        return proxy.ProxyClient.handleHeader(self, key, value)

    def handleResponsePart(self, buffer):
        """
        If we're not in streaming mode, buffer the response part(s).
        """
        if self._streaming:
            proxy.ProxyClient.handleResponsePart(self, buffer)
        else:
            # print "BUFFER!", buffer
            self._responsePartBuffer += buffer

    def handleResponseEnd(self):
        """
        If we're completing a chunked response, up-call to handle it like a
        regular reverse proxy.

        If we're completing a non-chunked response, fire the post-hooks.

        If we're completing a hijacked response, pass through the connection
        close.
        """
        if self.http:
            if self._streaming:
                return proxy.ProxyClient.handleResponseEnd(self)
            else:
                contentType = self.father.responseHeaders.getRawHeaders("content-type")
                if contentType:
                    contentType = contentType[0]
                else:
                    contentType = None
                body = self._responsePartBuffer
                self._fireListener(
                        {"PowerstripProtocolVersion": 1,
                         "ModifiedServerResponse":
                            {"Body": body,
                             "Code": self.father.code,
                             "ContentType": contentType}})
        else:
            self.father.transport.loseConnection()

    def rawDataReceived(self, data):
        if self.http:
            return proxy.ProxyClient.rawDataReceived(self, data)
        self.father.transport.write(data)

class DockerProxyClientFactory(proxy.ProxyClientFactory):
    protocol = DockerProxyClient
    _listener = None
    def onCreate(self, d):
        self._listener = d

    def _fireListener(self, result):
        if self._listener is not None:
            d = self._listener
            self._listener = None
            d.callback(result)

    def buildProtocol(self, addr):
        client = proxy.ProxyClientFactory.buildProtocol(self, addr)
        self._fireListener(client)
        return client

class DockerProxy(proxy.ReverseProxyResource):
    proxyClientFactoryClass = DockerProxyClientFactory

    def __init__(self, dockerAddr=None, dockerPort=None, dockerSocket=None,
            path='', reactor=reactor, config=None):
        """
        A docker proxy resource which knows how to connect to real Docker
        daemon either via socket (dockerSocket specified) or address + port for
        TCP connection (dockerAddr + dockerPort specified).
        """
        if config is None:
            # Try to get the configuration from the default place on the
            # filesystem.
            self.config = PluginConfiguration()
        else:
            self.config = config
        self.config.read_and_parse()
        self.parser = EndpointParser(self.config)
        Resource.__init__(self)
        self.host = dockerAddr
        self.port = dockerPort
        self.socket = dockerSocket
        self.path = path
        self.reactor = reactor
        proxy.ReverseProxyResource.__init__(self, dockerAddr, dockerPort, path, reactor) # NB dockerAddr is not actually used
        self.agent = Agent(reactor) # no connectionpool
        self.client = HTTPClient(self.agent)

    def render(self, request, reactor=reactor):
        # We are processing a leaf request.
        # Get the original request body from the client.
        skipPreHooks = False
        if request.requestHeaders.getRawHeaders('content-type') == ["application/json"]:
            originalRequestBody = request.content.read()
            request.content.seek(0) # hee hee
        elif request.requestHeaders.getRawHeaders('content-type') == ["application/tar"]:
            # We can't JSON encode binary data, so don't even try.
            skipPreHooks = True
            originalRequestBody = None
        else:
            originalRequestBody = None
        preHooks = []
        postHooks = []
        d = defer.succeed(None)
        for endpoint in self.parser.match_endpoint(request.method, request.uri.split("?")[0]):
            # It's possible for a request to match multiple endpoint
            # definitions.  Order of matched endpoint is not defined in
            # that case.
            adapters = self.config.endpoint(endpoint)
            preHooks.extend(adapters.pre)
            postHooks.extend(adapters.post)
        def callPreHook(result, hookURL):
            if result is None:
                newRequestBody = originalRequestBody
            else:
                newRequestBody = result["ModifiedClientRequest"]["Body"]
            return self.client.post(hookURL, json.dumps({
                        "PowerstripProtocolVersion": 1,
                        "Type": "pre-hook",
                        "ClientRequest": {
                            "Method": request.method,
                            "Request": request.uri,
                            "Body": newRequestBody,
                            # XXX This would need a ContentType header... if we
                            # were to support non-JSON POST bodies, like build
                            # contexts in pre-hooks.
                        }
                    }), headers={'Content-Type': ['application/json']})
        if not skipPreHooks:
            for preHook in preHooks:
                hookURL = self.config.adapter_uri(preHook)
                d.addCallback(callPreHook, hookURL=hookURL)
                d.addCallback(treq.json_content)
        def doneAllPrehooks(result):
            # Finally pass through the request to actual Docker.  For now we
            # mutate request in-place in such a way that ReverseProxyResource
            # understands it.
            if result is not None:
                requestBody = b""
                bodyFromAdapter = result["ModifiedClientRequest"]["Body"]
                if bodyFromAdapter is not None:
                    requestBody = bodyFromAdapter.encode("utf-8")
                request.content = StringIO.StringIO(requestBody)
                request.requestHeaders.setRawHeaders(b"content-length",
                        [str(len(requestBody))])
            ###########################
            # The following code is copied from t.w.proxy.ReverseProxy so that
            # clientFactory reference can be kept.
            if not self.socket:
                if self.port == 80:
                    host = self.host
                else:
                    host = "%s:%d" % (self.host, self.port)
                request.requestHeaders.setRawHeaders(b"host", [host])
            request.content.seek(0, 0)
            qs = urlparse.urlparse(request.uri)[4]
            if qs:
                rest = self.path + '?' + qs
            else:
                rest = self.path
            allRequestHeaders = request.getAllHeaders()
            if ("transfer-encoding" in allRequestHeaders
                    and allRequestHeaders["transfer-encoding"] == "chunked"):
                del allRequestHeaders["transfer-encoding"]
            # XXX Streaming the contents of the request body into memory could
            # cause OOM issues for large build contexts POSTed through
            # powerstrip.
            body = request.content.read()
            allRequestHeaders["content-length"] = str(len(body))
            clientFactory = self.proxyClientFactoryClass(
                request.method, rest, request.clientproto,
                allRequestHeaders, body, request)
            ###########################
            if self.socket:
                self.reactor.connectUNIX(self.socket, clientFactory)
            else:
                self.reactor.connectTCP(self.host, self.port, clientFactory)
            d = defer.Deferred()
            clientFactory.onCreate(d)
            return d
        d.addCallback(doneAllPrehooks)
        def inspect(client):
            # If there are no post-hooks, allow the response to be streamed
            # back to the client, rather than buffered.
            d = defer.Deferred()
            client.registerListener(d)
            if not postHooks:
                client.setStreamingMode(True)
            return d
        d.addCallback(inspect)
        def callPostHook(result, hookURL):
            serverResponse = result["ModifiedServerResponse"]
            return self.client.post(hookURL, json.dumps({
                        # TODO Write tests for the information provided to the adapter.
                        "PowerstripProtocolVersion": 1,
                        "Type": "post-hook",
                        "ClientRequest": {
                            "Method": request.method,
                            "Request": request.uri,
                            "Body": originalRequestBody,
                            },
                        "ServerResponse": {
                            "ContentType": serverResponse["ContentType"],
                            "Body": serverResponse["Body"],
                            "Code": serverResponse["Code"],
                        },
                    }), headers={'Content-Type': ['application/json']})
        # XXX Need to skip post-hooks for tar archives from e.g. docker export.
        for postHook in postHooks:
            hookURL = self.config.adapter_uri(postHook)
            d.addCallback(callPostHook, hookURL=hookURL)
            d.addCallback(treq.json_content)
        def sendFinalResponseToClient(result):
            # Write the final response to the client.
            request.write(result["ModifiedServerResponse"]["Body"].encode("utf-8"))
            request.finish()
        d.addCallback(sendFinalResponseToClient)
        def squashNoPostHooks(failure):
            failure.trap(NoPostHooks)
        d.addErrback(squashNoPostHooks)
        d.addErrback(log.err, 'while running chain')
        return NOT_DONE_YET

    def getChild(self, path, request):
        fragments = request.uri.split("/")
        fragments.pop(0)
        proxyArgs = (self.host, self.port, self.socket, self.path + '/' + urlquote(path, safe=""),
                     self.reactor)
        #if not request.postpath:
        resource = DockerProxy(*proxyArgs, config=self.config)
        return resource

class ServerProtocolFactory(server.Site):
    def __init__(self, dockerAddr=None, dockerPort=None, dockerSocket=None, config=None):
        self.root = DockerProxy(dockerAddr, dockerPort, dockerSocket, config=config)
        server.Site.__init__(self, self.root)
<|endoftext|>


<commit_before>#!/usr/bin/env python
"""
Cling Kernel for Jupyter

Talks to Cling via ctypes
"""

from __future__ import print_function

__version__ = '0.0.2'

import ctypes
from contextlib import contextmanager
from fcntl import fcntl, F_GETFL, F_SETFL
import os
import select
import sys
import threading

from traitlets import Unicode, Float
from ipykernel.kernelapp import IPKernelApp
from ipykernel.kernelbase import Kernel

libc = ctypes.CDLL(None)
try:
    c_stdout_p = ctypes.c_void_p.in_dll(libc, 'stdout')
except ValueError:
    # libc.stdout is has a funny name on OS X
    c_stdout_p = ctypes.c_void_p.in_dll(libc, '__stdoutp')

class ClingKernel(Kernel):
    """Cling Kernel for Jupyter"""
    implementation = 'cling_kernel'
    implementation_version = __version__
    language_version = 'X'

    banner = Unicode()
    def _banner_default(self):
        return 'cling-%s' % self.language_version
        return self._banner

        # codemirror_mode='clike' *should* work but doesn't, using the mimetype instead
    language_info = {'name': 'c++',
                     'codemirror_mode': 'text/x-c++src',
                     'mimetype': ' text/x-c++src',
                     'file_extension': '.c++'}

    cling = Unicode(config=True,
        help="Path to cling if not on your PATH."
    )
    flush_interval = Float(0.25, config=True)

    @contextmanager
    def forward_stdout(self):
        """Capture stdout and forward it as stream messages"""
        # create pipe for stdout
        stdout_fd = sys.__stdout__.fileno()
        save_stdout = os.dup(stdout_fd)
        pipe_out, pipe_in = os.pipe()
        os.dup2(pipe_in, stdout_fd)
        os.close(pipe_in)

        # make pipe_out non-blocking
        flags = fcntl(pipe_out, F_GETFL)
        fcntl(pipe_out, F_SETFL, flags|os.O_NONBLOCK)

        def forwarder(pipe, name='stdout'):
            """Forward bytes on a pipe to stream messages"""
            while True:
                r, w, x = select.select([pipe], [], [], self.flush_interval)
                if not r:
                    # nothing to read, flush libc's stdout and check again
                    libc.fflush(c_stdout_p)
                    continue
                data = os.read(pipe, 1024)
                if not data:
                    # pipe closed, we are done
                    break
                # send output
                self.session.send(self.iopub_socket, 'stream', {
                    'name': name,
                    'text': data.decode('utf8', 'replace'),
                }, parent=self._parent_header)

        t = threading.Thread(target=forwarder, args=(pipe_out,))
        t.start()
        try:
            yield
        finally:
            # flush the pipe
            libc.fflush(c_stdout_p)
            os.close(stdout_fd)
            t.join()

            # and restore original stdout
            os.close(pipe_out)
            os.dup2(save_stdout, stdout_fd)
            os.close(save_stdout)

    def run_cell(self, code, silent=False):
        """Dummy run cell while waiting for cling ctypes API"""
        import time

        for i in range(5):
            libc.printf((code + '\n').encode('utf8'))
            time.sleep(0.2 * i)

    def do_execute(self, code, silent, store_history=True,
                   user_expressions=None, allow_stdin=False):
        if not code.strip():
            return {
                'status': 'ok',
                'execution_count': self.execution_count,
                'payload': [],
                'user_expressions': {},
            }
        status = 'ok'

        with self.forward_stdout():
            self.run_cell(code, silent)

        reply = {
            'status': status,
            'execution_count': self.execution_count,
        }

        if status == 'error':
            err = {
                'ename': 'ename',
                'evalue': 'evalue',
                'traceback': [],
            }
            self.send_response(self.iopub_socket, 'error', err)
            reply.update(err)
        elif status == 'ok':
            reply.update({
                'payload': [],
                'user_expressions': {},
            })
        else:
            raise ValueError("Invalid status: %r" % status)

        return reply

class ClingKernelApp(IPKernelApp):
    kernel_class = ClingKernel
    def init_io(self):
        # disable io forwarding
        pass

def main():
    """launch a cling kernel"""
    ClingKernelApp.launch_instance()

if __name__ == '__main__':
    main()
<commit_msg>capture stderr as well<commit_after>#!/usr/bin/env python
"""
Cling Kernel for Jupyter

Talks to Cling via ctypes
"""

from __future__ import print_function

__version__ = '0.0.2'

import ctypes
from contextlib import contextmanager
from fcntl import fcntl, F_GETFL, F_SETFL
import os
import select
import sys
import threading

from traitlets import Unicode, Float
from ipykernel.kernelapp import IPKernelApp
from ipykernel.kernelbase import Kernel

libc = ctypes.CDLL(None)
try:
    c_stdout_p = ctypes.c_void_p.in_dll(libc, 'stdout')
    c_stderr_p = ctypes.c_void_p.in_dll(libc, 'stderr')
except ValueError:
    # libc.stdout is has a funny name on OS X
    c_stdout_p = ctypes.c_void_p.in_dll(libc, '__stdoutp')
    c_stderr_p = ctypes.c_void_p.in_dll(libc, '__stderrp')

class ClingKernel(Kernel):
    """Cling Kernel for Jupyter"""
    implementation = 'cling_kernel'
    implementation_version = __version__
    language_version = 'X'

    banner = Unicode()
    def _banner_default(self):
        return 'cling-%s' % self.language_version
        return self._banner

        # codemirror_mode='clike' *should* work but doesn't, using the mimetype instead
    language_info = {'name': 'c++',
                     'codemirror_mode': 'text/x-c++src',
                     'mimetype': ' text/x-c++src',
                     'file_extension': '.c++'}

    cling = Unicode(config=True,
        help="Path to cling if not on your PATH."
    )
    flush_interval = Float(0.25, config=True)

    @contextmanager
    def forward_stream(self, name):
        """Capture stdout and forward it as stream messages"""
        # create pipe for stdout
        if name == 'stdout':
            c_flush_p = c_stdout_p
        elif name == 'stderr':
            c_flush_p = c_stderr_p
        else:
            raise ValueError("Name must be stdout or stderr, not %r" % name)

        real_fd = getattr(sys, '__%s__' % name).fileno()
        save_fd = os.dup(real_fd)
        pipe_out, pipe_in = os.pipe()
        os.dup2(pipe_in, real_fd)
        os.close(pipe_in)

        # make pipe_out non-blocking
        flags = fcntl(pipe_out, F_GETFL)
        fcntl(pipe_out, F_SETFL, flags|os.O_NONBLOCK)

        def forwarder(pipe):
            """Forward bytes on a pipe to stream messages"""
            while True:
                r, w, x = select.select([pipe], [], [], self.flush_interval)
                if not r:
                    # nothing to read, flush libc's stdout and check again
                    libc.fflush(c_flush_p)
                    continue
                data = os.read(pipe, 1024)
                if not data:
                    # pipe closed, we are done
                    break
                # send output
                self.session.send(self.iopub_socket, 'stream', {
                    'name': name,
                    'text': data.decode('utf8', 'replace'),
                }, parent=self._parent_header)

        t = threading.Thread(target=forwarder, args=(pipe_out,))
        t.start()
        try:
            yield
        finally:
            # flush the pipe
            libc.fflush(c_flush_p)
            os.close(real_fd)
            t.join()

            # and restore original stdout
            os.close(pipe_out)
            os.dup2(save_fd, real_fd)
            os.close(save_fd)

    def run_cell(self, code, silent=False):
        """Dummy run cell while waiting for cling ctypes API"""
        import time

        for i in range(5):
            libc.printf((code + '\n').encode('utf8'))
            time.sleep(0.2 * i)

    def do_execute(self, code, silent, store_history=True,
                   user_expressions=None, allow_stdin=False):
        if not code.strip():
            return {
                'status': 'ok',
                'execution_count': self.execution_count,
                'payload': [],
                'user_expressions': {},
            }
        status = 'ok'

        with self.forward_stream('stdout'), self.forward_stream('stderr'):
            self.run_cell(code, silent)

        reply = {
            'status': status,
            'execution_count': self.execution_count,
        }

        if status == 'error':
            err = {
                'ename': 'ename',
                'evalue': 'evalue',
                'traceback': [],
            }
            self.send_response(self.iopub_socket, 'error', err)
            reply.update(err)
        elif status == 'ok':
            reply.update({
                'payload': [],
                'user_expressions': {},
            })
        else:
            raise ValueError("Invalid status: %r" % status)

        return reply

class ClingKernelApp(IPKernelApp):
    kernel_class = ClingKernel
    def init_io(self):
        # disable io forwarding
        pass

def main():
    """launch a cling kernel"""
    ClingKernelApp.launch_instance()

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


<commit_before>import pymel.core as pm
from PipelineTools.utilities import *
#reload(*)
@do_function_on('single')
def constraint_ctrl_to_loc(n):
    offgpname = ob.replace('_ctl','_offset_Gp')
    gpname = ob.replace('_ctl','_Gp')
    pm.group(n, n=offgpname)
    pm.group(offgpname, n=gpname)
    locname = ob.replace('_ctl','_loc')
    pm.pointConstraint(locname, gpname, o=(0,0,0), w=1)

@do_function_on('last')
def snap_to_near_vertex_midedge(ob_list,target):
    target_pos_list = [v.getPosition(space='world') for v in target.vtx]
    midedge_pos_list = [(e.getPoint(0, space='world')+e.getPoint(1, space='world'))/2 for e in target.e]
    print midedge_pos_list
    target_pos_list.extend(midedge_pos_list)
    #def get_closest_vert():
    for ob in ob_list:
        ob_pos = ob.getTranslation(space='world')
        distance = [(ob_pos.distanceTo(target_pos), target_pos) for target_pos in target_pos_list]
        distance.sort(key=lambda x:x[0])
        print ob,distance[0]
        ob.setTranslation(distance[0][1], space='world')
    #get_closest_vert()
    #return distance_list

@do_function_on('single')
def group_loc(ob):
    gpname = n + '_Gp'
    pm.group(n, n=gpname)

@do_function_on('single')
def make_loc_and_ctl_from_bon(ob):
    locname = ob.replace('_bon','_loc')
    loc = pm.spaceLocator(p=(0,0,0), n=locname)
    pc = pm.pointConstraint(n, loc, o=(0,0,0), w=1)
    pm.delete(pc[0])
    ctlname = ob.replace('_bon','_ctl')
    sph = pm.sphere(p=(0,0,0), ax=(0,1,0), ssw=0, esw=360, r=0.35, d=3, ut=False, tol=0.01, s=8, nsp=4, ch=False, n=ctlname)
    pc = pm.pointConstraint(n, sph, o=(0,0,0), w=1)
    pm.delete(pc[0])

def connect_Bs_control():
    faceBS = 'FaceBaseBS'
    browCnt = 'brow_ctl'
    eyeCnt = 'eye_ctl'
    mouthCnt = 'mouth_ctl'
    pm.connectAttr(browCnt+'.eyebrow_smile_L', faceBS+'.eyebrow_smile_L', f=True)
    pm.connectAttr(browCnt+'.eyebrow_smile_R', faceBS+'.eyebrow_smile_R', f=True)
    pm.connectAttr(browCnt+'.eyebrow_sad_L', faceBS+'.eyebrow_sad_L', f=True)
    pm.connectAttr(browCnt+'.eyebrow_sad_R', faceBS+'.eyebrow_sad_R', f=True)
    pm.connectAttr(browCnt+'.eyebrow_anger_L', faceBS+'.eyebrow_anger_L', f=True)
    pm.connectAttr(browCnt+'.eyebrow_anger_R', faceBS+'.eyebrow_anger_R', f=True)
    pm.connectAttr(eyeCnt+'.eye_close_L', faceBS+'.eye_close_L', f=True)
    pm.connectAttr(eyeCnt+'.eye_close_R', faceBS+'.eye_close_R', f=True)
    pm.connectAttr(eyeCnt+'.eye_smile_L', faceBS+'.eye_smile_L', f=True)
    pm.connectAttr(eyeCnt+'.eye_smile_R', faceBS+'.eye_smile_R', f=True)
    pm.connectAttr(eyeCnt+'.eye_anger_L', faceBS+'.eye_anger_L', f=True)
    pm.connectAttr(eyeCnt+'.eye_anger_R', faceBS+'.eye_anger_R', f=True)
    pm.connectAttr(eyeCnt+'.eye_open_L', faceBS+'.eye_open_L', f=True)
    pm.connectAttr(eyeCnt+'.eye_open_R', faceBS+'.eye_open_R', f=True)
    pm.connectAttr(mouthCnt+'.mouth_A', faceBS+'.mouth_A', f=True)
    pm.connectAttr(mouthCnt+'.mouth_I', faceBS+'.mouth_I', f=True)
    pm.connectAttr(mouthCnt+'.mouth_U', faceBS+'.mouth_U', f=True)
    pm.connectAttr(mouthCnt+'.mouth_E', faceBS+'.mouth_E', f=True)
    pm.connectAttr(mouthCnt+'.mouth_O', faceBS+'.mouth_O', f=True)
    pm.connectAttr(mouthCnt+'.mouth_shout', faceBS+'.mouth_shout', f=True)
    pm.connectAttr(mouthCnt+'.mouth_open', faceBS+'.mouth_open', f=True)
    pm.connectAttr(mouthCnt+'.mouth_smileclose', faceBS+'.mouth_smileclose', f=True)
    pm.connectAttr(mouthCnt+'.mouth_angerclose', faceBS+'.mouth_angerclose', f=True)

def connect_face_ctl():
    lC = 'lipCenter'
    nT = 'noseTop'
    ebL = 'eyebrowLeft'
    eL = 'eyeLeft'
    eSL = 'eyeSubLeft'
    nL = 'noseLeft'
    cL = 'cheekLeft'
    lL = 'lipLeft'
    ebR = 'eyebrowRight'
    eR = 'eyeRight'
    eSR = 'eyeSubRight'
    nR = 'noseRight'
    cR = 'cheekRight'
    lR = 'lipRight'
    for attr in ['translate', 'rotate', 'scale']:
        #Center
        pm.connectAttr(lC+'A_ctl.%s'%attr, lC+'A_bon.%s'%attr, f=True)
        pm.connectAttr(lC+'B_ctl.%s'%attr, lC+'B_bon.%s'%attr, f=True)
        pm.connectAttr(nT+'_ctl.%s'%attr, nT+'_bon.%s'%attr, f=True)
        #Left
        pm.connectAttr(ebL+'A_ctl.%s'%attr, ebL+'A_bon.%s'%attr, f=True)
        pm.connectAttr(ebL+'B_ctl.%s'%attr, ebL+'B_bon.%s'%attr, f=True)
        pm.connectAttr(ebL+'C_ctl.%s'%attr, ebL+'C_bon.%s'%attr, f=True)
        pm.connectAttr(eL+'A_ctl.%s'%attr, eL+'A_bon.%s'%attr, f=True)
        pm.connectAttr(eL+'B_ctl.%s'%attr, eL+'B_bon.%s'%attr, f=True)
        pm.connectAttr(eL+'C_ctl.%s'%attr, eL+'C_bon.%s'%attr, f=True)
        pm.connectAttr(eL+'D_ctl.%s'%attr, eL+'D_bon.%s'%attr, f=True)
        pm.connectAttr(eSL+'A_ctl.%s'%attr, eSL+'A_bon.%s'%attr, f=True)
        pm.connectAttr(eSL+'B_ctl.%s'%attr, eSL+'B_bon.%s'%attr, f=True)
        pm.connectAttr(eSL+'C_ctl.%s'%attr, eSL+'C_bon.%s'%attr, f=True)
        pm.connectAttr(eSL+'D_ctl.%s'%attr, eSL+'D_bon.%s'%attr, f=True)
        pm.connectAttr(nL+'A_ctl.%s'%attr, nL+'A_bon.%s'%attr, f=True)
        pm.connectAttr(cL+'A_ctl.%s'%attr, cL+'A_bon.%s'%attr, f=True)
        pm.connectAttr(cL+'B_ctl.%s'%attr, cL+'B_bon.%s'%attr, f=True)
        pm.connectAttr(cL+'C_ctl.%s'%attr, cL+'C_bon.%s'%attr, f=True)
        pm.connectAttr(lL+'A_ctl.%s'%attr, lL+'A_bon.%s'%attr, f=True)
        pm.connectAttr(lL+'B_ctl.%s'%attr, lL+'B_bon.%s'%attr, f=True)
        pm.connectAttr(lL+'C_ctl.%s'%attr, lL+'C_bon.%s'%attr, f=True)
        #Right
        pm.connectAttr(ebR+'A_ctl.%s'%attr, ebR+'A_bon.%s'%attr, f=True)
        pm.connectAttr(ebR+'B_ctl.%s'%attr, ebR+'B_bon.%s'%attr, f=True)
        pm.connectAttr(ebR+'C_ctl.%s'%attr, ebR+'C_bon.%s'%attr, f=True)
        pm.connectAttr(eR+'A_ctl.%s'%attr, eR+'A_bon.%s'%attr, f=True)
        pm.connectAttr(eR+'B_ctl.%s'%attr, eR+'B_bon.%s'%attr, f=True)
        pm.connectAttr(eR+'C_ctl.%s'%attr, eR+'C_bon.%s'%attr, f=True)
        pm.connectAttr(eR+'D_ctl.%s'%attr, eR+'D_bon.%s'%attr, f=True)
        pm.connectAttr(eSR+'A_ctl.%s'%attr, eSR+'A_bon.%s'%attr, f=True)
        pm.connectAttr(eSR+'B_ctl.%s'%attr, eSR+'B_bon.%s'%attr, f=True)
        pm.connectAttr(eSR+'C_ctl.%s'%attr, eSR+'C_bon.%s'%attr, f=True)
        pm.connectAttr(eSR+'D_ctl.%s'%attr, eSR+'D_bon.%s'%attr, f=True)
        pm.connectAttr(nR+'A_ctl.%s'%attr, nR+'A_bon.%s'%attr, f=True)
        pm.connectAttr(cR+'A_ctl.%s'%attr, cR+'A_bon.%s'%attr, f=True)
        pm.connectAttr(cR+'B_ctl.%s'%attr, cR+'B_bon.%s'%attr, f=True)
        pm.connectAttr(cR+'C_ctl.%s'%attr, cR+'C_bon.%s'%attr, f=True)
        pm.connectAttr(lR+'A_ctl.%s'%attr, lR+'A_bon.%s'%attr, f=True)
        pm.connectAttr(lR+'B_ctl.%s'%attr, lR+'B_bon.%s'%attr, f=True)
        pm.connectAttr(lR+'C_ctl.%s'%attr, lR+'C_bon.%s'%attr, f=True)

def connect_mouth_ctl():
    jaw = 'jaw'
    teeU = 'teethUpper'
    teeL = 'teethLower'
    tonR = 'tongueRight'
    tonC = 'tongueCenter'
    tonL = 'tongueLeft'
    tonM = 'tongueCenter'
    for attr in ['translate', 'rotate', 'scale']:
        pm.connectAttr(jaw+'_ctl.%s'%attr, jaw+'Root_bon.%s'%attr, f=True)
        pm.connectAttr(teeU+'_ctl.%s'%attr, teeU+'_bon.%s'%attr, f=True)
        pm.connectAttr(teeL+'_ctl.%s'%attr, teeL+'_bon.%s'%attr, f=True)
        pm.connectAttr(tonR+'A_ctl.%s'%attr, tonR+'A_bon.%s'%attr, f=True)
        pm.connectAttr(tonC+'A_ctl.%s'%attr, tonC+'A_bon.%s'%attr, f=True)
        pm.connectAttr(tonL+'A_ctl.%s'%attr, tonL+'A_bon.%s'%attr, f=True)
        pm.connectAttr(tonM+'ARoot_ctl.%s'%attr, tonM+'ARoot_bon.%s'%attr, f=True)
        pm.connectAttr(tonR+'B_ctl.%s'%attr, tonR+'B_bon.%s'%attr, f=True)
        pm.connectAttr(tonC+'B_ctl.%s'%attr, tonC+'B_bon.%s'%attr, f=True)
        pm.connectAttr(tonL+'B_ctl.%s'%attr, tonL+'B_bon.%s'%attr, f=True)
        pm.connectAttr(tonM+'BRoot_ctl.%s'%attr, tonM+'BRoot_bon.%s'%attr, f=True)
        pm.connectAttr(tonR+'C_ctl.%s'%attr, tonR+'C_bon.%s'%attr, f=True)
        pm.connectAttr(tonC+'C_ctl.%s'%attr, tonC+'C_bon.%s'%attr, f=True)
        pm.connectAttr(tonL+'C_ctl.%s'%attr, tonL+'C_bon.%s'%attr, f=True)
        pm.connectAttr(tonM+'CRoot_ctl.%s'%attr, tonM+'CRoot_bon.%s'%attr, f=True)
        pm.connectAttr(tonR+'D_ctl.%s'%attr, tonR+'D_bon.%s'%attr, f=True)
        pm.connectAttr(tonC+'D_ctl.%s'%attr, tonC+'D_bon.%s'%attr, f=True)
        pm.connectAttr(tonL+'D_ctl.%s'%attr, tonL+'D_bon.%s'%attr, f=True)
        pm.connectAttr(tonM+'DRoot_ctl.%s'%attr, tonM+'DRoot_bon.%s'%attr, f=True)<commit_msg>update<commit_after>import pymel.core as pm
from PipelineTools.utilities import *
#reload(*)
@do_function_on('single')
def constraint_ctrl_to_loc(n):
    offgpname = ob.replace('_ctl','_offset_Gp')
    gpname = ob.replace('_ctl','_Gp')
    pm.group(n, n=offgpname)
    pm.group(offgpname, n=gpname)
    locname = ob.replace('_ctl','_loc')
    pm.pointConstraint(locname, gpname, o=(0,0,0), w=1)
@timeit
@do_function_on('last')
def snap_to_near_vertex_midedge(ob_list,target):
    target_pos_list = [v.getPosition(space='world') for v in target.vtx]
    midedge_pos_list = [(e.getPoint(0, space='world')+e.getPoint(1, space='world'))/2 for e in target.e]
    print midedge_pos_list
    target_pos_list.extend(midedge_pos_list)
    #def get_closest_vert():
    for ob in ob_list:
        ob_pos = ob.getTranslation(space='world')
        distance = [(ob_pos.distanceTo(target_pos), target_pos) for target_pos in target_pos_list]
        distance.sort(key=lambda x:x[0])
        #print ob,distance[0]
        ob.setTranslation(distance[0][1], space='world')
    #get_closest_vert()
    #return distance_list

@do_function_on('single')
def group_loc(ob):
    gpname = n + '_Gp'
    pm.group(n, n=gpname)

@do_function_on('single')
def make_loc_and_ctl_from_bon(ob):
    locname = ob.replace('_bon','_loc')
    loc = pm.spaceLocator(p=(0,0,0), n=locname)
    pc = pm.pointConstraint(n, loc, o=(0,0,0), w=1)
    pm.delete(pc[0])
    ctlname = ob.replace('_bon','_ctl')
    sph = pm.sphere(p=(0,0,0), ax=(0,1,0), ssw=0, esw=360, r=0.35, d=3, ut=False, tol=0.01, s=8, nsp=4, ch=False, n=ctlname)
    pc = pm.pointConstraint(n, sph, o=(0,0,0), w=1)
    pm.delete(pc[0])

def connect_Bs_control():
    faceBS = 'FaceBaseBS'
    browCnt = 'brow_ctl'
    eyeCnt = 'eye_ctl'
    mouthCnt = 'mouth_ctl'
    pm.connectAttr(browCnt+'.eyebrow_smile_L', faceBS+'.eyebrow_smile_L', f=True)
    pm.connectAttr(browCnt+'.eyebrow_smile_R', faceBS+'.eyebrow_smile_R', f=True)
    pm.connectAttr(browCnt+'.eyebrow_sad_L', faceBS+'.eyebrow_sad_L', f=True)
    pm.connectAttr(browCnt+'.eyebrow_sad_R', faceBS+'.eyebrow_sad_R', f=True)
    pm.connectAttr(browCnt+'.eyebrow_anger_L', faceBS+'.eyebrow_anger_L', f=True)
    pm.connectAttr(browCnt+'.eyebrow_anger_R', faceBS+'.eyebrow_anger_R', f=True)
    pm.connectAttr(eyeCnt+'.eye_close_L', faceBS+'.eye_close_L', f=True)
    pm.connectAttr(eyeCnt+'.eye_close_R', faceBS+'.eye_close_R', f=True)
    pm.connectAttr(eyeCnt+'.eye_smile_L', faceBS+'.eye_smile_L', f=True)
    pm.connectAttr(eyeCnt+'.eye_smile_R', faceBS+'.eye_smile_R', f=True)
    pm.connectAttr(eyeCnt+'.eye_anger_L', faceBS+'.eye_anger_L', f=True)
    pm.connectAttr(eyeCnt+'.eye_anger_R', faceBS+'.eye_anger_R', f=True)
    pm.connectAttr(eyeCnt+'.eye_open_L', faceBS+'.eye_open_L', f=True)
    pm.connectAttr(eyeCnt+'.eye_open_R', faceBS+'.eye_open_R', f=True)
    pm.connectAttr(mouthCnt+'.mouth_A', faceBS+'.mouth_A', f=True)
    pm.connectAttr(mouthCnt+'.mouth_I', faceBS+'.mouth_I', f=True)
    pm.connectAttr(mouthCnt+'.mouth_U', faceBS+'.mouth_U', f=True)
    pm.connectAttr(mouthCnt+'.mouth_E', faceBS+'.mouth_E', f=True)
    pm.connectAttr(mouthCnt+'.mouth_O', faceBS+'.mouth_O', f=True)
    pm.connectAttr(mouthCnt+'.mouth_shout', faceBS+'.mouth_shout', f=True)
    pm.connectAttr(mouthCnt+'.mouth_open', faceBS+'.mouth_open', f=True)
    pm.connectAttr(mouthCnt+'.mouth_smileclose', faceBS+'.mouth_smileclose', f=True)
    pm.connectAttr(mouthCnt+'.mouth_angerclose', faceBS+'.mouth_angerclose', f=True)

def connect_face_ctl():
    lC = 'lipCenter'
    nT = 'noseTop'
    ebL = 'eyebrowLeft'
    eL = 'eyeLeft'
    eSL = 'eyeSubLeft'
    nL = 'noseLeft'
    cL = 'cheekLeft'
    lL = 'lipLeft'
    ebR = 'eyebrowRight'
    eR = 'eyeRight'
    eSR = 'eyeSubRight'
    nR = 'noseRight'
    cR = 'cheekRight'
    lR = 'lipRight'
    for attr in ['translate', 'rotate', 'scale']:
        #Center
        pm.connectAttr(lC+'A_ctl.%s'%attr, lC+'A_bon.%s'%attr, f=True)
        pm.connectAttr(lC+'B_ctl.%s'%attr, lC+'B_bon.%s'%attr, f=True)
        pm.connectAttr(nT+'_ctl.%s'%attr, nT+'_bon.%s'%attr, f=True)
        #Left
        pm.connectAttr(ebL+'A_ctl.%s'%attr, ebL+'A_bon.%s'%attr, f=True)
        pm.connectAttr(ebL+'B_ctl.%s'%attr, ebL+'B_bon.%s'%attr, f=True)
        pm.connectAttr(ebL+'C_ctl.%s'%attr, ebL+'C_bon.%s'%attr, f=True)
        pm.connectAttr(eL+'A_ctl.%s'%attr, eL+'A_bon.%s'%attr, f=True)
        pm.connectAttr(eL+'B_ctl.%s'%attr, eL+'B_bon.%s'%attr, f=True)
        pm.connectAttr(eL+'C_ctl.%s'%attr, eL+'C_bon.%s'%attr, f=True)
        pm.connectAttr(eL+'D_ctl.%s'%attr, eL+'D_bon.%s'%attr, f=True)
        pm.connectAttr(eSL+'A_ctl.%s'%attr, eSL+'A_bon.%s'%attr, f=True)
        pm.connectAttr(eSL+'B_ctl.%s'%attr, eSL+'B_bon.%s'%attr, f=True)
        pm.connectAttr(eSL+'C_ctl.%s'%attr, eSL+'C_bon.%s'%attr, f=True)
        pm.connectAttr(eSL+'D_ctl.%s'%attr, eSL+'D_bon.%s'%attr, f=True)
        pm.connectAttr(nL+'A_ctl.%s'%attr, nL+'A_bon.%s'%attr, f=True)
        pm.connectAttr(cL+'A_ctl.%s'%attr, cL+'A_bon.%s'%attr, f=True)
        pm.connectAttr(cL+'B_ctl.%s'%attr, cL+'B_bon.%s'%attr, f=True)
        pm.connectAttr(cL+'C_ctl.%s'%attr, cL+'C_bon.%s'%attr, f=True)
        pm.connectAttr(lL+'A_ctl.%s'%attr, lL+'A_bon.%s'%attr, f=True)
        pm.connectAttr(lL+'B_ctl.%s'%attr, lL+'B_bon.%s'%attr, f=True)
        pm.connectAttr(lL+'C_ctl.%s'%attr, lL+'C_bon.%s'%attr, f=True)
        #Right
        pm.connectAttr(ebR+'A_ctl.%s'%attr, ebR+'A_bon.%s'%attr, f=True)
        pm.connectAttr(ebR+'B_ctl.%s'%attr, ebR+'B_bon.%s'%attr, f=True)
        pm.connectAttr(ebR+'C_ctl.%s'%attr, ebR+'C_bon.%s'%attr, f=True)
        pm.connectAttr(eR+'A_ctl.%s'%attr, eR+'A_bon.%s'%attr, f=True)
        pm.connectAttr(eR+'B_ctl.%s'%attr, eR+'B_bon.%s'%attr, f=True)
        pm.connectAttr(eR+'C_ctl.%s'%attr, eR+'C_bon.%s'%attr, f=True)
        pm.connectAttr(eR+'D_ctl.%s'%attr, eR+'D_bon.%s'%attr, f=True)
        pm.connectAttr(eSR+'A_ctl.%s'%attr, eSR+'A_bon.%s'%attr, f=True)
        pm.connectAttr(eSR+'B_ctl.%s'%attr, eSR+'B_bon.%s'%attr, f=True)
        pm.connectAttr(eSR+'C_ctl.%s'%attr, eSR+'C_bon.%s'%attr, f=True)
        pm.connectAttr(eSR+'D_ctl.%s'%attr, eSR+'D_bon.%s'%attr, f=True)
        pm.connectAttr(nR+'A_ctl.%s'%attr, nR+'A_bon.%s'%attr, f=True)
        pm.connectAttr(cR+'A_ctl.%s'%attr, cR+'A_bon.%s'%attr, f=True)
        pm.connectAttr(cR+'B_ctl.%s'%attr, cR+'B_bon.%s'%attr, f=True)
        pm.connectAttr(cR+'C_ctl.%s'%attr, cR+'C_bon.%s'%attr, f=True)
        pm.connectAttr(lR+'A_ctl.%s'%attr, lR+'A_bon.%s'%attr, f=True)
        pm.connectAttr(lR+'B_ctl.%s'%attr, lR+'B_bon.%s'%attr, f=True)
        pm.connectAttr(lR+'C_ctl.%s'%attr, lR+'C_bon.%s'%attr, f=True)

def connect_mouth_ctl():
    jaw = 'jaw'
    teeU = 'teethUpper'
    teeL = 'teethLower'
    tonR = 'tongueRight'
    tonC = 'tongueCenter'
    tonL = 'tongueLeft'
    tonM = 'tongueCenter'
    for attr in ['translate', 'rotate', 'scale']:
        pm.connectAttr(jaw+'_ctl.%s'%attr, jaw+'Root_bon.%s'%attr, f=True)
        pm.connectAttr(teeU+'_ctl.%s'%attr, teeU+'_bon.%s'%attr, f=True)
        pm.connectAttr(teeL+'_ctl.%s'%attr, teeL+'_bon.%s'%attr, f=True)
        pm.connectAttr(tonR+'A_ctl.%s'%attr, tonR+'A_bon.%s'%attr, f=True)
        pm.connectAttr(tonC+'A_ctl.%s'%attr, tonC+'A_bon.%s'%attr, f=True)
        pm.connectAttr(tonL+'A_ctl.%s'%attr, tonL+'A_bon.%s'%attr, f=True)
        pm.connectAttr(tonM+'ARoot_ctl.%s'%attr, tonM+'ARoot_bon.%s'%attr, f=True)
        pm.connectAttr(tonR+'B_ctl.%s'%attr, tonR+'B_bon.%s'%attr, f=True)
        pm.connectAttr(tonC+'B_ctl.%s'%attr, tonC+'B_bon.%s'%attr, f=True)
        pm.connectAttr(tonL+'B_ctl.%s'%attr, tonL+'B_bon.%s'%attr, f=True)
        pm.connectAttr(tonM+'BRoot_ctl.%s'%attr, tonM+'BRoot_bon.%s'%attr, f=True)
        pm.connectAttr(tonR+'C_ctl.%s'%attr, tonR+'C_bon.%s'%attr, f=True)
        pm.connectAttr(tonC+'C_ctl.%s'%attr, tonC+'C_bon.%s'%attr, f=True)
        pm.connectAttr(tonL+'C_ctl.%s'%attr, tonL+'C_bon.%s'%attr, f=True)
        pm.connectAttr(tonM+'CRoot_ctl.%s'%attr, tonM+'CRoot_bon.%s'%attr, f=True)
        pm.connectAttr(tonR+'D_ctl.%s'%attr, tonR+'D_bon.%s'%attr, f=True)
        pm.connectAttr(tonC+'D_ctl.%s'%attr, tonC+'D_bon.%s'%attr, f=True)
        pm.connectAttr(tonL+'D_ctl.%s'%attr, tonL+'D_bon.%s'%attr, f=True)
        pm.connectAttr(tonM+'DRoot_ctl.%s'%attr, tonM+'DRoot_bon.%s'%attr, f=True)<|endoftext|>


<commit_before>import overpy

def getEmbassies():
	api = overpy.Overpass()
	# fetch all ways and nodes
	result = api.query("""
[out:json];
(
	(
		(
			node[amenity=embassy][country][target]["addr:street"]["addr:city"]["contact:phone"];
				- node[amenity=embassy][diplomatic][diplomatic!=embassy];
		);
			- node[amenity=embassy][!name][!"name:en"];
	);
	(
		(
			way[amenity=embassy][country][target]["addr:street"]["addr:city"]["contact:phone"];
				- way[amenity=embassy][diplomatic][diplomatic!=embassy];
		);
			- way[amenity=embassy][!name][!"name:en"];
	);
	(
		(
			rel[amenity=embassy][country][target]["addr:street"]["addr:city"]["contact:phone"];
				- rel[amenity=embassy][diplomatic][diplomatic!=embassy];
		);
			- rel[amenity=embassy][!name][!"name:en"];
	);
);
out tags;
		""")

	return result<commit_msg>Printing out embassy info<commit_after>import overpy

def getEmbassies():
	api = overpy.Overpass()
	# fetch all ways and nodes
	result = api.query("""
[out:json];
(
	(
		(
			node[amenity=embassy][country][target]["addr:street"]["addr:city"]["contact:phone"];
				- node[amenity=embassy][diplomatic][diplomatic!=embassy];
		);
			- node[amenity=embassy][!name][!"name:en"];
	);
	(
		(
			way[amenity=embassy][country][target]["addr:street"]["addr:city"]["contact:phone"];
				- way[amenity=embassy][diplomatic][diplomatic!=embassy];
		);
			- way[amenity=embassy][!name][!"name:en"];
	);
	(
		(
			rel[amenity=embassy][country][target]["addr:street"]["addr:city"]["contact:phone"];
				- rel[amenity=embassy][diplomatic][diplomatic!=embassy];
		);
			- rel[amenity=embassy][!name][!"name:en"];
	);
);
out tags;
		""")

	for l in ['ways', 'nodes', 'relations']:
		element = getattr(result, l)
		for e in element:
			embassy = e.tags
			print("{},{},{},{},{},{}".format(embassy['name'].encode("utf-8"),
				embassy['country'].encode("utf-8"),
				embassy['target'].encode("utf-8"),
				embassy['addr:street'].encode("utf-8"),
				embassy['addr:city'].encode("utf-8"),
				embassy['contact:phone'].encode("utf-8")))

	return result<|endoftext|>