<commit_before>#! /usr/bin/python
# -*- coding: utf-8 -*-

import numpy as np
import tensorflow as tf
import tensorlayer as tl

from tensorlayer.layers.core import Layer
# from tensorlayer.layers.core import LayersConfig

from tensorlayer import logging

__all__ = [
    'OneHot',
    'Word2vecEmbedding',
    'Embedding',
    'AverageEmbedding',
]

class OneHot(Layer):
    """
    The :class:`OneHot` class is the starting layer of a neural network, see ``tf.one_hot``.

    Parameters
    ----------
    depth : None or int
        If the input indices is rank N, the output will have rank N+1. The new axis is created at dimension `axis` (default: the new axis is appended at the end).
    on_value : None or number
        The value to represnt `ON`. If None, it will default to the value 1.
    off_value : None or number
        The value to represnt `OFF`. If None, it will default to the value 0.
    axis : None or int
        The axis.
    dtype : None or TensorFlow dtype
        The data type, None means tf.float32.
    name : str
        A unique layer name.

    Examples
    ---------
    >>> import tensorflow as tf
    >>> import tensorlayer as tl
    >>> x = tf.placeholder(tf.int32, shape=[None])
    >>> net = tl.layers.OneHot(x, depth=8)
    (?, 8)

    """

    def __init__(self, depth=None, on_value=None, off_value=None, axis=None, dtype=None, name=None):  #'input'):

        # super(OneHotInput, self).__init__(prev_layer=inputs, name=name)
        super().__init__(name)
        self.depth = depth
        self.on_value = on_value
        self.off_value = off_value
        self.axis = axis
        self.dtype = dtype
        logging.info("OneHotInput  %s: %s" % (self.name, str(inputs.shape.as_list())))

        if self.depth is None:
            raise RuntimeError(self.__class__.__name__ + ": depth == None the number of output units is undefined")

    def build(self, inputs):
        pass

    def forward(self, inputs):
        """
        Parameters
        ----------
        inputs : input tensor
            The input of a network.
        """
        outputs = tf.one_hot(
            inputs, self.depth, on_value=self.on_value, off_value=self.off_value, axis=self.axis, dtype=self.dtype
        )
        return outputs

class Word2vecEmbedding(Layer):
    """
    The :class:`Word2vecEmbedding` class is a fully connected layer.
    For Word Embedding, words are input as integer index.
    The output is the embedded word vector.

    Parameters
    ----------
    vocabulary_size : int
        The size of vocabulary, number of words
    embedding_size : int
        The number of embedding dimensions
    num_sampled : int
        The mumber of negative examples for NCE loss
    nce_loss_args : dictionary
        The arguments for tf.nn.nce_loss()
    E_init : initializer
        The initializer for initializing the embedding matrix
    nce_W_init : initializer
        The initializer for initializing the nce decoder weight matrix
    nce_b_init : initializer
        The initializer for initializing of the nce decoder bias vector
    name : str
        A unique layer name

    Attributes
    ----------
    outputs : Tensor
        The embedding layer outputs.
    normalized_embeddings : Tensor
        Normalized embedding matrix.
    nce_weights : Tensor
        The NCE weights
    nce_biases: Tensor
        The NCE biases

    Examples
    --------
    With TensorLayer (Example in `examples/text_word_embedding/tutorial_word2vec_basic.py`)

    >>> import tensorflow as tf
    >>> import tensorlayer as tl
    >>> batch_size = 8
    >>> train_inputs = tf.placeholder(tf.int32, shape=[batch_size])
    >>> train_labels = tf.placeholder(tf.int32, shape=[batch_size, 1])

    >>> net_in = tl.layers.Input([batch_size], dtype=tf.int32)
    >>> emb_net = tl.layers.Word2vecEmbedding(
    >>>       vocabulary_size=vocabulary_size,
    >>>       embedding_size=embedding_size,
    >>>       num_sampled=num_sampled,
    >>>       nce_loss_args={},
    >>>       E_init=tl.initializers.random_uniform(minval=-1.0, maxval=1.0),
    >>>       nce_W_init=tl.initializers.truncated_normal(stddev=float(1.0 / np.sqrt(embedding_size))),
    >>>       nce_b_init=tl.initializers.constant(value=0.0),
    >>>       name='word2vec_layer',
    >>> )(net_in)

    >>> model = tl.models.Model(inputs=net_in, outputs=emb_net, name="word2vec_model")

    >>> nce_cost = tf.reduce_mean(
    >>>     input_tensor=tf.nn.nce_loss(
    >>>         weights=emb_net.nce_weights,
    >>>         biases=emb_net.nce_biases,
    >>>         inputs=model(train_inputs, is_train=True),
    >>>         labels=train_labels,  #self.train_labels,
    >>>         num_sampled=emb_net.num_sampled,
    >>>         num_classes=emb_net.vocabulary_size,
    >>>         **emb_net.nce_loss_args
    >>>    )
    >>> )

    >>> train_params = model.weights
    >>> train_op = tf.train.GradientDescentOptimizer(learning_rate).minimize(nce_cost, var_list=train_params)

    >>> normalized_embeddings = emb_net.normalized_embeddings

    References
    ----------
    `tensorflow/examples/tutorials/word2vec/word2vec_basic.py <https://github.com/tensorflow/tensorflow/blob/r0.7/tensorflow/examples/tutorials/word2vec/word2vec_basic.py>`__

    """

    def __init__(
            self,
            vocabulary_size=80000,
            embedding_size=200,
            num_sampled=64,
            nce_loss_args=None,
            E_init=tl.initializers.random_uniform(minval=-1.0, maxval=1.0),
            nce_W_init=tl.initializers.truncated_normal(stddev=0.03),
            nce_b_init=tl.initializers.constant(value=0.0),
            name=None,  #'word2vec',
    ):

        super(Word2vecEmbedding, self).__init__(name)
        self.vocabulary_size = vocabulary_size
        self.embedding_size = embedding_size
        self.num_sampled = num_sampled
        self.nce_loss_args = nce_loss_args
        self.E_init = E_init
        self.nce_W_init = nce_W_init
        self.nce_b_init = nce_b_init
        logging.info("Word2vecEmbedding %s: (%d, %d)" % (self.name, self.vocabulary_size, self.embedding_size))

    def build(self, inputs):
        # Look up embeddings for inputs.
        # Note: a row of 'embeddings' is the vector representation of a word.
        # for the sake of speed, it is better to slice the embedding matrix
        # instead of transfering a word id to one-hot-format vector and then
        # multiply by the embedding matrix.
        # embed is the outputs of the hidden layer (embedding layer), it is a
        # row vector with 'embedding_size' values.

        self.embeddings = self._get_weights(
            "embeddings", shape=(self.vocabulary_size, self.embedding_size), init=self.E_init,
        )

        self.normalized_embeddings = tf.nn.l2_normalize(self.embeddings, 1)

        # Construct the variables for the NCE loss (i.e. negative sampling)
        self.nce_weights = self._get_weights(
            "nce_weights", shape=(self.vocabulary_size, self.embedding_size), init=self.nce_W_init,
        )

        self.nce_biases = self._get_weights(
            "nce_biases", shape=(self.vocabulary_size,), init=self.nce_b_init,
        )

    def forward(self, inputs):
        """
        Parameters
        ----------
        inputs : input tensor
            The input of a network. For word inputs, please use integer index format, 2D tensor : [batch_size, num_steps(num_words)]
        train_labels : input tensor
            For word labels. integer index format
        """
        outputs = tf.nn.embedding_lookup(params=self.embeddings, ids=inputs)

        return outputs

class Embedding(Layer):
    """
    The :class:`Embedding` class is a look-up table for word embedding.

    Word content are accessed using integer indexes, then the output is the embedded word vector.
    To train a word embedding matrix, you can used :class:`Word2vecEmbedding`.
    If you have a pre-trained matrix, you can assign the parameters into it.

    Parameters
    ----------
    vocabulary_size : int
        The size of vocabulary, number of words.
    embedding_size : int
        The number of embedding dimensions.
    E_init : initializer
        The initializer for the embedding matrix.
    E_init_args : dictionary
        The arguments for embedding matrix initializer.
    name : str
        A unique layer name.

    Attributes
    ----------
    outputs : tensor
        The embedding layer output is a 3D tensor in the shape: (batch_size, num_steps(num_words), embedding_size).

    Examples
    --------
    >>> import tensorflow as tf
    >>> import tensorlayer as tl
    >>> batch_size = 8
    >>> x = tf.placeholder(tf.int32, shape=(batch_size, ))
    >>> net = tl.layers.Embedding(inputs=x, vocabulary_size=1000, embedding_size=50, name='embed')
    (8, 50)

    """

    def __init__(
            self,
            # inputs,
            vocabulary_size=80000,
            embedding_size=200,
            E_init=tf.compat.v1.initializers.random_uniform(-0.1, 0.1),
            E_init_args=None,
            name=None,  #'embedding',
    ):
        # super(Embedding, self).__init__(prev_layer=inputs, E_init_args=E_init_args, name=name)
        super().__init__(name)
        self.vocabulary_size = vocabulary_size
        self.embedding_size = embedding_size
        self.E_init = E_init
        self.E_init_args = E_init_args

        logging.info("Embedding %s: (%d, %d)" % (self.name, self.vocabulary_size, self.embedding_size))

    def build(self, inputs):
        # self.embeddings = tf.compat.v1.get_variable(
        #     name=self.name + '/embeddings', shape=(self.vocabulary_size, self.embedding_size), initializer=self.E_init,
        #     dtype=LayersConfig.tf_dtype, **self.E_init_args
        # )
        self.embeddings = self._get_weights(
            "embeddings", shape=(self.vocabulary_size, self.embedding_size), init=self.E_init,
            init_args=self.E_init_args
        )

    def forward(self, inputs):
        """
        Parameters
        ----------
        inputs : placeholder
            The input of a network. For word inputs.
            Please use integer index format, 2D tensor : (batch_size, num_steps(num_words)).
        """
        outputs = tf.nn.embedding_lookup(params=self.embeddings, ids=inputs)
        return outputs

class AverageEmbedding(Layer):
    """The :class:`AverageEmbedding` averages over embeddings of inputs.
    This is often used as the input layer for models like DAN[1] and FastText[2].

    Parameters
    ----------
    vocabulary_size : int
        The size of vocabulary.
    embedding_size : int
        The dimension of the embedding vectors.
    pad_value : int
        The scalar padding value used in inputs, 0 as default.
    E_init : initializer
        The initializer of the embedding matrix.
    name : str
        A unique layer name.

    References
    ----------
    - [1] Iyyer, M., Manjunatha, V., Boyd-Graber, J., & Daum’e III, H. (2015). Deep Unordered Composition Rivals Syntactic Methods for Text Classification. In Association for Computational Linguistics.
    - [2] Joulin, A., Grave, E., Bojanowski, P., & Mikolov, T. (2016). `Bag of Tricks for Efficient Text Classification. <http://arxiv.org/abs/1607.01759>`__

    Examples
    ---------
    >>> import tensorflow as tf
    >>> import tensorlayer as tl
    >>> batch_size = 8
    >>> length = 5
    >>> net = tl.layers.Input([batch_size, length], dtype=tf.int32)
    >>> net = tl.layers.AverageEmbedding(vocabulary_size=1000, embedding_size=50, name='avg')(net)
    (8, 50)

    """

    def __init__(
            self,
            vocabulary_size,
            embedding_size,
            pad_value=0,
            E_init=tl.initializers.random_uniform(-0.1, 0.1),
            name=None,  # 'average_embedding',
    ):

        # super(AverageEmbedding, self).__init__(prev_layer=inputs, embeddings_kwargs=embeddings_kwargs, name=name)
        super(AverageEmbedding, self).__init__(name)
        self.vocabulary_size = vocabulary_size
        self.embedding_size = embedding_size
        self.pad_value = pad_value
        self.E_init = E_init
        logging.info("AverageEmbedding %s: (%d, %d)" % (self.name, self.vocabulary_size, self.embedding_size))

        # if embeddings_kwargs is None:
        #     embeddings_kwargs = {}

    def build(self, inputs_shape):
        if len(inputs_shape) != 2:
            raise ValueError('inputs must be of size batch_size * batch_sentence_length')

        # self.embeddings = tf.compat.v1.get_variable(
        #     name=self.name + '/embeddings', shape=(self.vocabulary_size, self.embedding_size),
        #     initializer=self.E_init, dtype=LayersConfig.tf_dtype, **self.E_init_argst
        # )
        self.embeddings = self._get_weights(
            "embeddings", shape=(self.vocabulary_size, self.embedding_size), init=self.E_init,
        )

    def forward(self, inputs):
        """
        Parameters
        ----------
        inputs : input tensor
            The network input.
            For word inputs, please use integer index format, 2D tensor: (batch_size, num_steps(num_words)).
        """
        word_embeddings = tf.nn.embedding_lookup(
            params=self.embeddings,
            ids=inputs,
            name='word_embeddings',
        )

        # Zero out embeddings of pad value
        masks = tf.not_equal(inputs, self.pad_value, name='masks')
        word_embeddings *= tf.cast(
            tf.expand_dims(masks, axis=-1),
            dtype=tf.float32
        )
        sum_word_embeddings = tf.reduce_sum(input_tensor=word_embeddings, axis=1)

        # Count number of non-padding words in each sentence
        sentence_lengths = tf.math.count_nonzero(
            masks,
            axis=1,
            keepdims=True,
            dtype=tf.float32,
            name='sentence_lengths',
        )

        sentence_embeddings = tf.divide(
            sum_word_embeddings,
            sentence_lengths + 1e-8,  # Add epsilon to avoid dividing by 0
            name='sentence_embeddings'
        )

        outputs = sentence_embeddings

        return outputs
<commit_msg>embedding refactored, documentation updated<commit_after>#! /usr/bin/python
# -*- coding: utf-8 -*-

import numpy as np
import tensorflow as tf
import tensorlayer as tl

from tensorlayer.layers.core import Layer
# from tensorlayer.layers.core import LayersConfig

from tensorlayer import logging

__all__ = [
    'OneHot',
    'Word2vecEmbedding',
    'Embedding',
    'AverageEmbedding',
]

class OneHot(Layer):
    """
    The :class:`OneHot` class is the starting layer of a neural network, see ``tf.one_hot``.

    Parameters
    ----------
    depth : None or int
        If the input indices is rank N, the output will have rank N+1. The new axis is created at dimension `axis` (default: the new axis is appended at the end).
    on_value : None or number
        The value to represnt `ON`. If None, it will default to the value 1.
    off_value : None or number
        The value to represnt `OFF`. If None, it will default to the value 0.
    axis : None or int
        The axis.
    dtype : None or TensorFlow dtype
        The data type, None means tf.float32.
    name : str
        A unique layer name.

    Examples
    ---------
    >>> import tensorflow as tf
    >>> import tensorlayer as tl
    >>> x = tf.placeholder(tf.int32, shape=[None])
    >>> net = tl.layers.OneHot(x, depth=8)
    (?, 8)

    """

    def __init__(self, depth=None, on_value=None, off_value=None, axis=None, dtype=None, name=None):  #'input'):

        # super(OneHotInput, self).__init__(prev_layer=inputs, name=name)
        super().__init__(name)
        self.depth = depth
        self.on_value = on_value
        self.off_value = off_value
        self.axis = axis
        self.dtype = dtype
        logging.info("OneHotInput  %s: %s" % (self.name, str(inputs.shape.as_list())))

        if self.depth is None:
            raise RuntimeError(self.__class__.__name__ + ": depth == None the number of output units is undefined")

    def build(self, inputs_shape):
        pass

    def forward(self, inputs):
        """
        Parameters
        ----------
        inputs : input tensor
            The input of a network.
        """
        outputs = tf.one_hot(
            inputs, self.depth, on_value=self.on_value, off_value=self.off_value, axis=self.axis, dtype=self.dtype
        )
        return outputs

class Word2vecEmbedding(Layer):
    """
    The :class:`Word2vecEmbedding` class is a fully connected layer.
    For Word Embedding, words are input as integer index.
    The output is the embedded word vector.

    Parameters
    ----------
    vocabulary_size : int
        The size of vocabulary, number of words
    embedding_size : int
        The number of embedding dimensions
    num_sampled : int
        The mumber of negative examples for NCE loss
    nce_loss_args : dictionary
        The arguments for tf.nn.nce_loss()
    E_init : initializer
        The initializer for initializing the embedding matrix
    nce_W_init : initializer
        The initializer for initializing the nce decoder weight matrix
    nce_b_init : initializer
        The initializer for initializing of the nce decoder bias vector
    name : str
        A unique layer name

    Attributes
    ----------
    outputs : Tensor
        The embedding layer outputs.
    normalized_embeddings : Tensor
        Normalized embedding matrix.
    nce_weights : Tensor
        The NCE weights
    nce_biases: Tensor
        The NCE biases

    Examples
    --------
    Word2Vec With TensorLayer (Example in `examples/text_word_embedding/tutorial_word2vec_basic.py`)

    >>> import tensorflow as tf
    >>> import tensorlayer as tl
    >>> batch_size = 8
    >>> train_inputs = tf.placeholder(tf.int32, shape=[batch_size])
    >>> train_labels = tf.placeholder(tf.int32, shape=[batch_size, 1])

    >>> net_in = tl.layers.Input([batch_size], dtype=tf.int32)
    >>> emb_net = tl.layers.Word2vecEmbedding(
    >>>       vocabulary_size=vocabulary_size,
    >>>       embedding_size=embedding_size,
    >>>       num_sampled=num_sampled,
    >>>       nce_loss_args={},
    >>>       E_init=tl.initializers.random_uniform(minval=-1.0, maxval=1.0),
    >>>       nce_W_init=tl.initializers.truncated_normal(stddev=float(1.0 / np.sqrt(embedding_size))),
    >>>       nce_b_init=tl.initializers.constant(value=0.0),
    >>>       name='word2vec_layer',
    >>> )(net_in)

    >>> model = tl.models.Model(inputs=net_in, outputs=emb_net, name="word2vec_model")

    >>> nce_cost = tf.reduce_mean(
    >>>     input_tensor=tf.nn.nce_loss(
    >>>         weights=emb_net.nce_weights,
    >>>         biases=emb_net.nce_biases,
    >>>         inputs=model(train_inputs, is_train=True),
    >>>         labels=train_labels,  #self.train_labels,
    >>>         num_sampled=emb_net.num_sampled,
    >>>         num_classes=emb_net.vocabulary_size,
    >>>         **emb_net.nce_loss_args
    >>>    )
    >>> )

    >>> train_params = model.weights
    >>> train_op = tf.train.GradientDescentOptimizer(learning_rate).minimize(nce_cost, var_list=train_params)

    >>> normalized_embeddings = emb_net.normalized_embeddings

    References
    ----------
    `tensorflow/examples/tutorials/word2vec/word2vec_basic.py <https://github.com/tensorflow/tensorflow/blob/r0.7/tensorflow/examples/tutorials/word2vec/word2vec_basic.py>`__

    """

    def __init__(
            self,
            vocabulary_size=80000,
            embedding_size=200,
            num_sampled=64,
            nce_loss_args=None,
            E_init=tl.initializers.random_uniform(minval=-1.0, maxval=1.0),
            nce_W_init=tl.initializers.truncated_normal(stddev=0.03),
            nce_b_init=tl.initializers.constant(value=0.0),
            name=None,  #'word2vec',
    ):

        super(Word2vecEmbedding, self).__init__(name)
        self.vocabulary_size = vocabulary_size
        self.embedding_size = embedding_size
        self.num_sampled = num_sampled
        self.nce_loss_args = nce_loss_args
        self.E_init = E_init
        self.nce_W_init = nce_W_init
        self.nce_b_init = nce_b_init
        logging.info("Word2vecEmbedding %s: (%d, %d)" % (self.name, self.vocabulary_size, self.embedding_size))

    def build(self, inputs_shape):
        """
        Parameters
        ----------
        inputs_shape : tuple
            the shape of inputs tensor
        """
        # Look up embeddings for inputs.
        # Note: a row of 'embeddings' is the vector representation of a word.
        # for the sake of speed, it is better to slice the embedding matrix
        # instead of transferring a word id to one-hot-format vector and then
        # multiply by the embedding matrix.
        # embed is the outputs of the hidden layer (embedding layer), it is a
        # row vector with 'embedding_size' values.

        self.embeddings = self._get_weights(
            "embeddings", shape=(self.vocabulary_size, self.embedding_size), init=self.E_init,
        )

        self.normalized_embeddings = tf.nn.l2_normalize(self.embeddings, 1)

        # Construct the variables for the NCE loss (i.e. negative sampling)
        self.nce_weights = self._get_weights(
            "nce_weights", shape=(self.vocabulary_size, self.embedding_size), init=self.nce_W_init,
        )

        self.nce_biases = self._get_weights(
            "nce_biases", shape=(self.vocabulary_size,), init=self.nce_b_init,
        )

    def forward(self, inputs):
        """
        Parameters
        ----------
        inputs : input tensor
            The input of a network
        """
        outputs = tf.nn.embedding_lookup(params=self.embeddings, ids=inputs)

        return outputs

class Embedding(Layer):
    """
    The :class:`Embedding` class is a look-up table for word embedding.

    Word content are accessed using integer indexes, then the output is the embedded word vector.
    To train a word embedding matrix, you can used :class:`Word2vecEmbedding`.
    If you have a pre-trained matrix, you can assign the parameters into it.

    Parameters
    ----------
    vocabulary_size : int
        The size of vocabulary, number of words.
    embedding_size : int
        The number of embedding dimensions.
    E_init : initializer
        The initializer for the embedding matrix.
    E_init_args : dictionary
        The arguments for embedding matrix initializer.
    name : str
        A unique layer name.

    Attributes
    ----------
    outputs : tensor
        The embedding layer output is a 3D tensor in the shape: (batch_size, num_steps(num_words), embedding_size).

    Examples
    --------
    >>> import tensorflow as tf
    >>> import tensorlayer as tl
    >>> batch_size = 8
    >>> net = tl.layers.Input([batch_size, 100], dtype=tf.int32)
    >>> net = tl.layers.Embedding(vocabulary_size=1000, embedding_size=50, name='embed')(net)
    >>> print(net.outputs)
    <tf.Tensor 'embedding_lookup/Identity:0' shape=(8, 100, 50) dtype=float32>

    """

    def __init__(
            self,
            vocabulary_size=80000,
            embedding_size=200,
            E_init=tl.initializers.random_uniform(-0.1, 0.1),
            name=None,  #'embedding',
    ):
        super(Embedding, self).__init__(name)
        self.vocabulary_size = vocabulary_size
        self.embedding_size = embedding_size
        self.E_init = E_init

        logging.info("Embedding %s: (%d, %d)" % (self.name, self.vocabulary_size, self.embedding_size))

    def build(self, inputs_shape):
        """
        Parameters
        ----------
        inputs_shape : tuple
            the shape of inputs tensor
        """

        self.embeddings = self._get_weights(
            "embeddings", shape=(self.vocabulary_size, self.embedding_size), init=self.E_init,
        )

    def forward(self, inputs):
        """
        Parameters
        ----------
        inputs : Tensor
            The input of a network.
        """
        outputs = tf.nn.embedding_lookup(params=self.embeddings, ids=inputs)
        return outputs

class AverageEmbedding(Layer):
    """The :class:`AverageEmbedding` averages over embeddings of inputs.
    This is often used as the input layer for models like DAN[1] and FastText[2].

    Parameters
    ----------
    vocabulary_size : int
        The size of vocabulary.
    embedding_size : int
        The dimension of the embedding vectors.
    pad_value : int
        The scalar padding value used in inputs, 0 as default.
    E_init : initializer
        The initializer of the embedding matrix.
    name : str
        A unique layer name.

    Attributes
    ----------
    outputs : tensor
        The embedding layer output is a 2D tensor in the shape: (batch_size, embedding_size).

    References
    ----------
    - [1] Iyyer, M., Manjunatha, V., Boyd-Graber, J., & Daum’e III, H. (2015). Deep Unordered Composition Rivals Syntactic Methods for Text Classification. In Association for Computational Linguistics.
    - [2] Joulin, A., Grave, E., Bojanowski, P., & Mikolov, T. (2016). `Bag of Tricks for Efficient Text Classification. <http://arxiv.org/abs/1607.01759>`__

    Examples
    ---------
    >>> import tensorflow as tf
    >>> import tensorlayer as tl
    >>> batch_size = 8
    >>> length = 5
    >>> net = tl.layers.Input([batch_size, length], dtype=tf.int32)
    >>> net = tl.layers.AverageEmbedding(vocabulary_size=1000, embedding_size=50, name='avg')(net)
    >>> print(net.outputs)
    <tf.Tensor 'sentence_embeddings:0' shape=(8, 50) dtype=float32>

    """

    def __init__(
            self,
            vocabulary_size,
            embedding_size,
            pad_value=0,
            E_init=tl.initializers.random_uniform(-0.1, 0.1),
            name=None,  # 'average_embedding',
    ):

        super(AverageEmbedding, self).__init__(name)
        self.vocabulary_size = vocabulary_size
        self.embedding_size = embedding_size
        self.pad_value = pad_value
        self.E_init = E_init
        logging.info("AverageEmbedding %s: (%d, %d)" % (self.name, self.vocabulary_size, self.embedding_size))

    def build(self, inputs_shape):
        """
        Parameters
        ----------
        inputs_shape : tuple
            the shape of inputs tensor.
        """
        if len(inputs_shape) != 2:
            raise ValueError('inputs must be of size (batch_size, sentence_length)')

        self.embeddings = self._get_weights(
            "embeddings", shape=(self.vocabulary_size, self.embedding_size), init=self.E_init,
        )

    def forward(self, inputs):
        """
        Parameters
        ----------
        inputs : tensor
            The network input.
            For word inputs, please use integer index format, 2D tensor: (batch_size, sentence_length).
        """
        word_embeddings = tf.nn.embedding_lookup(
            params=self.embeddings,
            ids=inputs,
            name='word_embeddings',
        )

        # Zero out embeddings of pad value
        masks = tf.not_equal(inputs, self.pad_value, name='masks')
        word_embeddings *= tf.cast(
            tf.expand_dims(masks, axis=-1),
            dtype=tf.float32
        )
        sum_word_embeddings = tf.reduce_sum(input_tensor=word_embeddings, axis=1)

        # Count number of non-padding words in each sentence
        sentence_lengths = tf.math.count_nonzero(
            masks,
            axis=1,
            keepdims=True,
            dtype=tf.float32,
            name='sentence_lengths',
        )

        sentence_embeddings = tf.divide(
            sum_word_embeddings,
            sentence_lengths + 1e-8,  # Add epsilon to avoid dividing by 0
            name='sentence_embeddings'
        )

        outputs = sentence_embeddings

        return outputs
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# File: batch_norm.py
# Author: Yuxin Wu <ppwwyyxx@gmail.com>

import tensorflow as tf
from tensorflow.contrib.framework import add_model_variable
from tensorflow.python.training import moving_averages

from ..tfutils.tower import get_current_tower_context
from ..utils import logger
from .common import layer_register

__all__ = ['BatchNorm', 'BatchRenorm']

# decay: being too close to 1 leads to slow start-up. torch use 0.9.
# eps: torch: 1e-5. Lasagne: 1e-4

# XXX This is deprecated. Only kept for future reference.
@layer_register(log_shape=False)
def BatchNormV1(x, use_local_stat=None, decay=0.9, epsilon=1e-5):
    shape = x.get_shape().as_list()
    assert len(shape) in [2, 4]

    n_out = shape[-1]  # channel
    assert n_out is not None
    beta = tf.get_variable('beta', [n_out],
                           initializer=tf.constant_initializer())
    gamma = tf.get_variable('gamma', [n_out],
                            initializer=tf.constant_initializer(1.0))

    if len(shape) == 2:
        batch_mean, batch_var = tf.nn.moments(x, [0], keep_dims=False)
    else:
        batch_mean, batch_var = tf.nn.moments(x, [0, 1, 2], keep_dims=False)
    # just to make a clear name.
    batch_mean = tf.identity(batch_mean, 'mean')
    batch_var = tf.identity(batch_var, 'variance')

    emaname = 'EMA'
    ctx = get_current_tower_context()
    if use_local_stat is None:
        use_local_stat = ctx.is_training
    if use_local_stat != ctx.is_training:
        logger.warn("[BatchNorm] use_local_stat != is_training")

    if use_local_stat:
        # training tower
        if ctx.is_training:
            # reuse = tf.get_variable_scope().reuse
            with tf.variable_scope(tf.get_variable_scope(), reuse=False):
                # BatchNorm in reuse scope can be tricky! Moving mean/variance are not reused
                with tf.name_scope(None):  # https://github.com/tensorflow/tensorflow/issues/2740
                    # if reuse=True, try to find and use the existing statistics
                    # how to use multiple tensors to update one EMA? seems impossbile
                    ema = tf.train.ExponentialMovingAverage(decay=decay, name=emaname)
                    ema_apply_op = ema.apply([batch_mean, batch_var])
                    ema_mean, ema_var = ema.average(batch_mean), ema.average(batch_var)
                    if ctx.is_main_training_tower:
                        # inside main training tower
                        add_model_variable(ema_mean)
                        add_model_variable(ema_var)
    else:
        # no apply() is called here, no magic vars will get created,
        # no reuse issue will happen
        assert not ctx.is_training
        with tf.name_scope(None):
            ema = tf.train.ExponentialMovingAverage(decay=decay, name=emaname)
            mean_var_name = ema.average_name(batch_mean)
            var_var_name = ema.average_name(batch_var)
            if ctx.is_main_tower:
                # main tower, but needs to use global stat. global stat must be from outside
                # when reuse=True, the desired variable name could
                # actually be different, because a different var is created
                # for different reuse tower
                ema_mean = tf.get_variable('mean/' + emaname, [n_out])
                ema_var = tf.get_variable('variance/' + emaname, [n_out])
            else:
                # use statistics in another tower
                G = tf.get_default_graph()
                ema_mean = ctx.find_tensor_in_main_tower(G, mean_var_name + ':0')
                ema_var = ctx.find_tensor_in_main_tower(G, var_var_name + ':0')

    if use_local_stat:
        batch = tf.cast(tf.shape(x)[0], tf.float32)
        mul = tf.where(tf.equal(batch, 1.0), 1.0, batch / (batch - 1))
        batch_var = batch_var * mul  # use unbiased variance estimator in training

        with tf.control_dependencies([ema_apply_op] if ctx.is_training else []):
            # only apply EMA op if is_training
            return tf.nn.batch_normalization(
                x, batch_mean, batch_var, beta, gamma, epsilon, 'output')
    else:
        return tf.nn.batch_normalization(
            x, ema_mean, ema_var, beta, gamma, epsilon, 'output')

def get_bn_variables(n_out, use_scale, use_bias, gamma_init):
    if use_bias:
        beta = tf.get_variable('beta', [n_out], initializer=tf.constant_initializer())
    else:
        beta = tf.zeros([n_out], name='beta')
    if use_scale:
        gamma = tf.get_variable('gamma', [n_out], initializer=gamma_init)
    else:
        gamma = tf.ones([n_out], name='gamma')
    # x * gamma + beta

    moving_mean = tf.get_variable('mean/EMA', [n_out],
                                  initializer=tf.constant_initializer(), trainable=False)
    moving_var = tf.get_variable('variance/EMA', [n_out],
                                 initializer=tf.constant_initializer(), trainable=False)
    return beta, gamma, moving_mean, moving_var

def update_bn_ema(xn, batch_mean, batch_var, moving_mean, moving_var, decay):
    # TODO is there a way to use zero_debias in multi-GPU?
    update_op1 = moving_averages.assign_moving_average(
        moving_mean, batch_mean, decay, zero_debias=False,
        name='mean_ema_op')
    update_op2 = moving_averages.assign_moving_average(
        moving_var, batch_var, decay, zero_debias=False,
        name='var_ema_op')
    add_model_variable(moving_mean)
    add_model_variable(moving_var)

    # seems faster than delayed update, but might behave otherwise in distributed settings.
    with tf.control_dependencies([update_op1, update_op2]):
        return tf.identity(xn, name='output')

@layer_register(log_shape=False)
def BatchNorm(x, use_local_stat=None, decay=0.9, epsilon=1e-5,
              use_scale=True, use_bias=True,
              gamma_init=tf.constant_initializer(1.0), data_format='NHWC'):
    """
    Batch Normalization layer, as described in the paper:
    `Batch Normalization: Accelerating Deep Network Training by
    Reducing Internal Covariance Shift <http://arxiv.org/abs/1502.03167>`_.

    Args:
        x (tf.Tensor): a 4D or 2D tensor. When 4D, the layout should match data_format.
        use_local_stat (bool): whether to use mean/var of the current batch or the moving average.
            Defaults to True in training and False in inference.
        decay (float): decay rate of moving average.
        epsilon (float): epsilon to avoid divide-by-zero.
        use_scale, use_bias (bool): whether to use the extra affine transformation or not.
        gamma_init: initializer for gamma (the scale).

    Returns:
        tf.Tensor: a tensor named ``output`` with the same shape of x.

    Variable Names:

    * ``beta``: the bias term. Will be zero-inited by default.
    * ``gamma``: the scale term. Will be one-inited by default.
        Input will be transformed by ``x * gamma + beta``.
    * ``mean/EMA``: the moving average of mean.
    * ``variance/EMA``: the moving average of variance.

    Note:
        In multi-GPU training, moving averages across GPUs are not aggregated.
        This is consistent with most frameworks.

        However, all GPUs use the moving averages on the first GPU (instead of
        their own), this is inconsistent with most frameworks (but consistent
        with the official inceptionv3 example).
    """
    shape = x.get_shape().as_list()
    assert len(shape) in [2, 4]
    if len(shape) == 2:
        data_format = 'NHWC'    # error using NCHW? (see #190)
    if data_format == 'NCHW':
        n_out = shape[1]
    else:
        n_out = shape[-1]  # channel
    if len(shape) == 2:
        x = tf.reshape(x, [-1, 1, 1, n_out])
    assert n_out is not None, "Input to BatchNorm cannot have unknown channels!"

    beta, gamma, moving_mean, moving_var = get_bn_variables(n_out, use_scale, use_bias, gamma_init)

    ctx = get_current_tower_context()
    if use_local_stat is None:
        use_local_stat = ctx.is_training
    if use_local_stat != ctx.is_training:
        # we allow the use of local_stat in testing (only print warnings)
        # because it is useful to certain applications.
        logger.warn("[BatchNorm] use_local_stat != is_training")

    if use_local_stat:
        xn, batch_mean, batch_var = tf.nn.fused_batch_norm(
            x, gamma, beta, epsilon=epsilon,
            is_training=True, data_format=data_format)
    else:
        assert not ctx.is_training, "In training, local statistics has to be used!"
        if data_format == 'NCHW':
            # fused is slower in inference, but support NCHW
            xn, _, _ = tf.nn.fused_batch_norm(
                x, gamma, beta,
                moving_mean, moving_var,
                epsilon=epsilon, is_training=False, data_format=data_format)
        else:
            xn = tf.nn.batch_normalization(     # work only for NHWC when moving_mean is a vector
                x, moving_mean, moving_var, beta, gamma, epsilon)

    if len(shape) == 2:
        axis = [2, 3] if data_format == 'NCHW' else [1, 2]
        xn = tf.squeeze(xn, axis)

    # maintain EMA only on one GPU.
    if ctx.is_main_training_tower:
        return update_bn_ema(xn, batch_mean, batch_var, moving_mean, moving_var, decay)
    else:
        return tf.identity(xn, name='output')

# TODO support NCHW
@layer_register(log_shape=False)
def BatchRenorm(x, rmax, dmax, decay=0.9, epsilon=1e-5,
                use_scale=True, use_bias=True):
    """
    Batch Renormalization layer, as described in the paper:
    `Batch Renormalization: Towards Reducing Minibatch Dependence in Batch-Normalized Models
    <https://arxiv.org/abs/1702.03275>`_.

    Args:
        x (tf.Tensor): a NHWC or NC tensor.
        rmax, dmax (tf.Tensor): a scalar tensor, the maximum allowed corrections.
        decay (float): decay rate of moving average.
        epsilon (float): epsilon to avoid divide-by-zero.
        use_scale, use_bias (bool): whether to use the extra affine transformation or not.

    Returns:
        tf.Tensor: a tensor named ``output`` with the same shape of x.

    Variable Names:

    * ``beta``: the bias term.
    * ``gamma``: the scale term. Input will be transformed by ``x * gamma + beta``.
    * ``mean/EMA``: the moving average of mean.
    * ``variance/EMA``: the moving average of variance.
    """

    shape = x.get_shape().as_list()
    assert len(shape) in [2, 4]
    n_out = shape[-1]
    if len(shape) == 2:
        x = tf.reshape(x, [-1, 1, 1, n_out])
    beta, gamma, moving_mean, moving_var = get_bn_variables(
        n_out, use_scale, use_bias, tf.constant_initializer(1.0))

    ctx = get_current_tower_context()
    use_local_stat = ctx.is_training
    # for BatchRenorm, use_local_stat should always be is_training, unless a
    # different usage comes out in the future.

    if use_local_stat:
        xn, batch_mean, batch_var = tf.nn.fused_batch_norm(x, gamma, beta,
                                                           epsilon=epsilon, is_training=True)
        inv_sigma = tf.rsqrt(moving_var, 'inv_sigma')
        r = tf.stop_gradient(tf.clip_by_value(
            tf.sqrt(batch_var) * inv_sigma, 1.0 / rmax, rmax))
        d = tf.stop_gradient(tf.clip_by_value(
            (batch_mean - moving_mean) * inv_sigma,
            -dmax, dmax))
        xn = xn * r + d
    else:
        xn = tf.nn.batch_normalization(
            x, moving_mean, moving_var, beta, gamma, epsilon)

    if len(shape) == 2:
        xn = tf.squeeze(xn, [1, 2])
    if ctx.is_main_training_tower:
        return update_bn_ema(xn, batch_mean, batch_var, moving_mean, moving_var, decay)
    else:
        return tf.identity(xn, name='output')
<commit_msg>always use non-fused op for BatchNorm inference; support NCHW for BatchRenorm<commit_after>#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# File: batch_norm.py
# Author: Yuxin Wu <ppwwyyxx@gmail.com>

import tensorflow as tf
from tensorflow.contrib.framework import add_model_variable
from tensorflow.python.training import moving_averages

from ..tfutils.tower import get_current_tower_context
from ..utils import logger
from .common import layer_register

__all__ = ['BatchNorm', 'BatchRenorm']

# decay: being too close to 1 leads to slow start-up. torch use 0.9.
# eps: torch: 1e-5. Lasagne: 1e-4

# XXX This is deprecated. Only kept for future reference.
@layer_register(log_shape=False)
def BatchNormV1(x, use_local_stat=None, decay=0.9, epsilon=1e-5):
    shape = x.get_shape().as_list()
    assert len(shape) in [2, 4]

    n_out = shape[-1]  # channel
    assert n_out is not None
    beta = tf.get_variable('beta', [n_out],
                           initializer=tf.constant_initializer())
    gamma = tf.get_variable('gamma', [n_out],
                            initializer=tf.constant_initializer(1.0))

    if len(shape) == 2:
        batch_mean, batch_var = tf.nn.moments(x, [0], keep_dims=False)
    else:
        batch_mean, batch_var = tf.nn.moments(x, [0, 1, 2], keep_dims=False)
    # just to make a clear name.
    batch_mean = tf.identity(batch_mean, 'mean')
    batch_var = tf.identity(batch_var, 'variance')

    emaname = 'EMA'
    ctx = get_current_tower_context()
    if use_local_stat is None:
        use_local_stat = ctx.is_training
    if use_local_stat != ctx.is_training:
        logger.warn("[BatchNorm] use_local_stat != is_training")

    if use_local_stat:
        # training tower
        if ctx.is_training:
            # reuse = tf.get_variable_scope().reuse
            with tf.variable_scope(tf.get_variable_scope(), reuse=False):
                # BatchNorm in reuse scope can be tricky! Moving mean/variance are not reused
                with tf.name_scope(None):  # https://github.com/tensorflow/tensorflow/issues/2740
                    # if reuse=True, try to find and use the existing statistics
                    # how to use multiple tensors to update one EMA? seems impossbile
                    ema = tf.train.ExponentialMovingAverage(decay=decay, name=emaname)
                    ema_apply_op = ema.apply([batch_mean, batch_var])
                    ema_mean, ema_var = ema.average(batch_mean), ema.average(batch_var)
                    if ctx.is_main_training_tower:
                        # inside main training tower
                        add_model_variable(ema_mean)
                        add_model_variable(ema_var)
    else:
        # no apply() is called here, no magic vars will get created,
        # no reuse issue will happen
        assert not ctx.is_training
        with tf.name_scope(None):
            ema = tf.train.ExponentialMovingAverage(decay=decay, name=emaname)
            mean_var_name = ema.average_name(batch_mean)
            var_var_name = ema.average_name(batch_var)
            if ctx.is_main_tower:
                # main tower, but needs to use global stat. global stat must be from outside
                # when reuse=True, the desired variable name could
                # actually be different, because a different var is created
                # for different reuse tower
                ema_mean = tf.get_variable('mean/' + emaname, [n_out])
                ema_var = tf.get_variable('variance/' + emaname, [n_out])
            else:
                # use statistics in another tower
                G = tf.get_default_graph()
                ema_mean = ctx.find_tensor_in_main_tower(G, mean_var_name + ':0')
                ema_var = ctx.find_tensor_in_main_tower(G, var_var_name + ':0')

    if use_local_stat:
        batch = tf.cast(tf.shape(x)[0], tf.float32)
        mul = tf.where(tf.equal(batch, 1.0), 1.0, batch / (batch - 1))
        batch_var = batch_var * mul  # use unbiased variance estimator in training

        with tf.control_dependencies([ema_apply_op] if ctx.is_training else []):
            # only apply EMA op if is_training
            return tf.nn.batch_normalization(
                x, batch_mean, batch_var, beta, gamma, epsilon, 'output')
    else:
        return tf.nn.batch_normalization(
            x, ema_mean, ema_var, beta, gamma, epsilon, 'output')

def get_bn_variables(n_out, use_scale, use_bias, gamma_init):
    if use_bias:
        beta = tf.get_variable('beta', [n_out], initializer=tf.constant_initializer())
    else:
        beta = tf.zeros([n_out], name='beta')
    if use_scale:
        gamma = tf.get_variable('gamma', [n_out], initializer=gamma_init)
    else:
        gamma = tf.ones([n_out], name='gamma')
    # x * gamma + beta

    moving_mean = tf.get_variable('mean/EMA', [n_out],
                                  initializer=tf.constant_initializer(), trainable=False)
    moving_var = tf.get_variable('variance/EMA', [n_out],
                                 initializer=tf.constant_initializer(), trainable=False)
    return beta, gamma, moving_mean, moving_var

def update_bn_ema(xn, batch_mean, batch_var, moving_mean, moving_var, decay):
    # TODO is there a way to use zero_debias in multi-GPU?
    update_op1 = moving_averages.assign_moving_average(
        moving_mean, batch_mean, decay, zero_debias=False,
        name='mean_ema_op')
    update_op2 = moving_averages.assign_moving_average(
        moving_var, batch_var, decay, zero_debias=False,
        name='var_ema_op')
    add_model_variable(moving_mean)
    add_model_variable(moving_var)

    # seems faster than delayed update, but might behave otherwise in distributed settings.
    with tf.control_dependencies([update_op1, update_op2]):
        return tf.identity(xn, name='output')

def reshape_for_bn(param, ndims, chan, data_format):
    if ndims == 2:
        shape = [1, chan]
    else:
        shape = [1, 1, 1, chan] if data_format == 'NHWC' else [1, chan, 1, 1]
    return tf.reshape(param, shape)

@layer_register(log_shape=False)
def BatchNorm(x, use_local_stat=None, decay=0.9, epsilon=1e-5,
              use_scale=True, use_bias=True,
              gamma_init=tf.constant_initializer(1.0), data_format='NHWC'):
    """
    Batch Normalization layer, as described in the paper:
    `Batch Normalization: Accelerating Deep Network Training by
    Reducing Internal Covariance Shift <http://arxiv.org/abs/1502.03167>`_.

    Args:
        x (tf.Tensor): a 4D or 2D tensor. When 4D, the layout should match data_format.
        use_local_stat (bool): whether to use mean/var of the current batch or the moving average.
            Defaults to True in training and False in inference.
        decay (float): decay rate of moving average.
        epsilon (float): epsilon to avoid divide-by-zero.
        use_scale, use_bias (bool): whether to use the extra affine transformation or not.
        gamma_init: initializer for gamma (the scale).

    Returns:
        tf.Tensor: a tensor named ``output`` with the same shape of x.

    Variable Names:

    * ``beta``: the bias term. Will be zero-inited by default.
    * ``gamma``: the scale term. Will be one-inited by default.
        Input will be transformed by ``x * gamma + beta``.
    * ``mean/EMA``: the moving average of mean.
    * ``variance/EMA``: the moving average of variance.

    Note:
        In multi-GPU training, moving averages across GPUs are not aggregated.
        This is consistent with most frameworks.

        However, all GPUs use the moving averages on the first GPU (instead of
        their own), this is inconsistent with most frameworks (but consistent
        with the official inceptionv3 example).
    """
    shape = x.get_shape().as_list()
    ndims = len(shape)
    assert ndims in [2, 4]
    if ndims == 2:
        data_format = 'NHWC'
    if data_format == 'NCHW':
        n_out = shape[1]
    else:
        n_out = shape[-1]  # channel
    assert n_out is not None, "Input to BatchNorm cannot have unknown channels!"
    beta, gamma, moving_mean, moving_var = get_bn_variables(n_out, use_scale, use_bias, gamma_init)

    ctx = get_current_tower_context()
    if use_local_stat is None:
        use_local_stat = ctx.is_training
    elif use_local_stat != ctx.is_training:
        # we allow the use of local_stat in testing (only print warnings)
        # because it is useful to certain applications.
        logger.warn("[BatchNorm] use_local_stat != is_training")

    if use_local_stat:
        if ndims == 2:
            x = tf.reshape(x, [-1, 1, 1, n_out])    # fused_bn only takes 4D input
            # fused_bn has error using NCHW? (see #190)

        xn, batch_mean, batch_var = tf.nn.fused_batch_norm(
            x, gamma, beta, epsilon=epsilon,
            is_training=True, data_format=data_format)

        if ndims == 2:
            xn = tf.squeeze(xn, [1, 2])
    else:
        assert not ctx.is_training, "In training, local statistics has to be used!"
        # non-fused op is faster for inference
        if ndims == 4 and data_format == 'NCHW':
            [g, b, mm, mv] = [reshape_for_bn(_, ndims, n_out, data_format)
                              for _ in [gamma, beta, moving_mean, moving_var]]
            xn = tf.nn.batch_normalization(x, mm, mv, b, g, epsilon)
        else:
            # avoid the reshape if possible (when channel is the last dimension)
            xn = tf.nn.batch_normalization(
                x, moving_mean, moving_var, beta, gamma, epsilon)

    # maintain EMA only on one GPU.
    if ctx.is_main_training_tower:
        return update_bn_ema(xn, batch_mean, batch_var, moving_mean, moving_var, decay)
    else:
        return tf.identity(xn, name='output')

# TODO support NCHW
@layer_register(log_shape=False)
def BatchRenorm(x, rmax, dmax, decay=0.9, epsilon=1e-5,
                use_scale=True, use_bias=True, data_format='NHWC'):
    """
    Batch Renormalization layer, as described in the paper:
    `Batch Renormalization: Towards Reducing Minibatch Dependence in Batch-Normalized Models
    <https://arxiv.org/abs/1702.03275>`_.

    Args:
        x (tf.Tensor): a NHWC or NC tensor.
        rmax, dmax (tf.Tensor): a scalar tensor, the maximum allowed corrections.
        decay (float): decay rate of moving average.
        epsilon (float): epsilon to avoid divide-by-zero.
        use_scale, use_bias (bool): whether to use the extra affine transformation or not.

    Returns:
        tf.Tensor: a tensor named ``output`` with the same shape of x.

    Variable Names:

    * ``beta``: the bias term.
    * ``gamma``: the scale term. Input will be transformed by ``x * gamma + beta``.
    * ``mean/EMA``: the moving average of mean.
    * ``variance/EMA``: the moving average of variance.
    """

    shape = x.get_shape().as_list()
    ndims = len(shape)
    assert ndims in [2, 4]
    if ndims == 2:
        data_format = 'NHWC'    # error using NCHW? (see #190)
    if data_format == 'NCHW':
        n_out = shape[1]
    else:
        n_out = shape[-1]  # channel
    assert n_out is not None, "Input to BatchRenorm cannot have unknown channels!"

    beta, gamma, moving_mean, moving_var = get_bn_variables(
        n_out, use_scale, use_bias, tf.constant_initializer(1.0))

    ctx = get_current_tower_context()
    use_local_stat = ctx.is_training
    # for BatchRenorm, use_local_stat should always be is_training, unless a
    # different usage comes out in the future.

    if use_local_stat:
        if ndims == 2:
            x = tf.reshape(x, [-1, 1, 1, n_out])

        xn, batch_mean, batch_var = tf.nn.fused_batch_norm(
            x, gamma, beta, epsilon=epsilon, is_training=True, data_format=data_format)

        inv_sigma = tf.rsqrt(moving_var, 'inv_sigma')
        r = tf.stop_gradient(tf.clip_by_value(
            tf.sqrt(batch_var) * inv_sigma, 1.0 / rmax, rmax))
        d = tf.stop_gradient(tf.clip_by_value(
            (batch_mean - moving_mean) * inv_sigma,
            -dmax, dmax))
        r = reshape_for_bn(r, ndims, n_out, data_format)
        d = reshape_for_bn(d, ndims, n_out, data_format)
        xn = xn * r + d

        if ndims == 2:
            xn = tf.squeeze(xn, [1, 2])

    else:
        if ndims == 4 and data_format == 'NCHW':
            [g, b, mm, mv] = [reshape_for_bn(_, ndims, n_out, data_format)
                              for _ in [gamma, beta, moving_mean, moving_var]]
            xn = tf.nn.batch_normalization(x, mm, mv, b, g, epsilon)
        else:
            xn = tf.nn.batch_normalization(
                x, moving_mean, moving_var, beta, gamma, epsilon)

    if ctx.is_main_training_tower:
        return update_bn_ema(xn, batch_mean, batch_var, moving_mean, moving_var, decay)
    else:
        return tf.identity(xn, name='output')
<|endoftext|>


<commit_before>"""
Wrapper class for all ternary plotting functions.
"""

from collections import namedtuple
from functools import partial

import numpy
from matplotlib import pyplot

from . import heatmapping
from . import lines
from . import plotting
from .helpers import project_point, convert_coordinates_sequence

BackgroundParameters = namedtuple('BackgroundParameters', ['color', 'alpha', 'zorder'])

def figure(ax=None, scale=None, permutation=None):
    """
    Wraps a Matplotlib AxesSubplot or generates a new one. Emulates matplotlib's
    > figure, ax = pyplot.subplots()

    Parameters
    ----------
    ax: AxesSubplot, None
        The matplotlib AxesSubplot to wrap
    scale: float, None
        The scale factor of the ternary plot
    """

    ternary_ax = TernaryAxesSubplot(ax=ax, scale=scale, permutation=permutation)
    return ternary_ax.get_figure(), ternary_ax

def mpl_redraw_callback(event, tax):
    """
    Callback to properly rotate and redraw text labels when the plot is drawn
    or resized.

    Parameters
    ----------
    event: a matplotlib event
        either 'resize_event' or 'draw_event'
    tax: TernaryAxesSubplot
         the TernaryAxesSubplot
    """
    tax._redraw_labels()

class TernaryAxesSubplot(object):
    """
    Wrapper for python-ternary and matplotlib figure. Parameters for member
    functions simply pass through to ternary's functions with the same names.
    This class manages the matplotlib axes, the scale, and the boundary scale
    to ease the use of ternary plotting functions.
    """

    def __init__(self, ax=None, scale=None, permutation=None):
        if not scale:
            scale = 1.0
        if ax:
            self.ax = ax
        else:
            _, self.ax = pyplot.subplots()
        self.set_scale(scale=scale)
        self._permutation = permutation
        self._boundary_scale = scale
        # Container for the axis labels supplied by the user
        self._labels = dict()
        self._corner_labels = dict()
        self._ticks = dict()
        # Container for the redrawing of labels
        self._to_remove = []
        self._connect_callbacks()
        # Background
        self._background_parameters = None
        # Cache for the background triangle object, so it can be removed and redrawn as needed.
        self._background_triangle = None
        self.set_background_color(color="whitesmoke", zorder=-1000, alpha=0.75)

    def _connect_callbacks(self):
        """Connect resize matplotlib callbacks."""
        figure = self.get_figure()
        callback = partial(mpl_redraw_callback, tax=self)
        event_names = ('resize_event', 'draw_event')
        for event_name in event_names:
            figure.canvas.mpl_connect(event_name, callback)

    def __repr__(self):
        return "TernaryAxesSubplot: %s" % self.ax.__hash__()

    def get_axes(self):
        """Returns the underlying matplotlib AxesSubplot object."""
        return self.ax

    def get_figure(self):
        """Return the underlying matplotlib figure object."""
        ax = self.get_axes()
        return ax.get_figure()

    def set_scale(self, scale=None):
        self._scale = scale
        self.resize_drawing_canvas()

    def get_scale(self):
        return self._scale

    def set_axis_limits(self, axis_limits=None):
        """
        Set min and max data limits for each of the three axes.

        axis_limits = dict
            keys are 'b','l' and 'r' for the three axes
            vals are lists of the min and max values for the axis in
            data units.
        """
        self._axis_limits = axis_limits

    def get_axis_limits(self):
        return self._axis_limits

    # Title and Axis Labels

    def set_title(self, title, **kwargs):
        """Sets the title on the underlying matplotlib AxesSubplot."""
        ax = self.get_axes()
        ax.set_title(title, **kwargs)

    def left_axis_label(self, label, position=None,  rotation=60, offset=0.08,
                        **kwargs):
        """
        Sets the label on the left axis.

        Parameters
        ----------
        label: String
            The axis label
        position: 3-Tuple of floats, None
            The position of the text label
        rotation: float, 60
            The angle of rotation of the label
        offset: float,
            Used to compute the distance of the label from the axis
        kwargs:
            Any kwargs to pass through to matplotlib.
        """

        if not position:
            position = (-offset, 3./5, 2./5)
        self._labels["left"] = (label, position, rotation, kwargs)

    def right_axis_label(self, label, position=None, rotation=-60, offset=0.08,
                         **kwargs):

        """
        Sets the label on the right axis.

        Parameters
        ----------
        label: String
            The axis label
        position: 3-Tuple of floats, None
            The position of the text label
        rotation: float, -60
            The angle of rotation of the label
        offset: float,
            Used to compute the distance of the label from the axis
        kwargs:
            Any kwargs to pass through to matplotlib.
        """

        if not position:
            position = (2. / 5 + offset, 3. / 5, 0)
        self._labels["right"] = (label, position, rotation, kwargs)

    def bottom_axis_label(self, label, position=None, rotation=0, offset=0.02,
                          **kwargs):
        """
        Sets the label on the bottom axis.

        Parameters
        ----------
        label: String
            The axis label
        position: 3-Tuple of floats, None
            The position of the text label
        rotation: float, 0
            The angle of rotation of the label
        offset: float,
            Used to compute the distance of the label from the axis
        kwargs:
            Any kwargs to pass through to matplotlib.
        """

        if not position:
            position = (0.5, -offset / 2., 0.5)
        self._labels["bottom"] = (label, position, rotation, kwargs)

    def right_corner_label(self, label, position=None, rotation=0, offset=0.08,
                           **kwargs):
        """
        Sets the label on the right corner (complements left axis).

        Parameters
        ----------
        label: String
            The axis label
        position: 3-Tuple of floats, None
            The position of the text label
        rotation: float, 0
            The angle of rotation of the label
        offset: float,
            Used to compute the distance of the label from the axis
        kwargs:
            Any kwargs to pass through to matplotlib.
        """

        if not position:
            position = (1, offset / 2, 0)
        self._corner_labels["right"] = (label, position, rotation, kwargs)

    def left_corner_label(self, label, position=None, rotation=0, offset=0.08,
                          **kwargs):
        """
        Sets the label on the left corner (complements right axis.)

        Parameters
        ----------
        label: string
            The axis label
        position: 3-Tuple of floats, None
            The position of the text label
        rotation: float, 0
            The angle of rotation of the label
        offset: float,
            Used to compute the distance of the label from the axis
        kwargs:
            Any kwargs to pass through to matplotlib.
        """

        if not position:
            position = (-offset / 2, offset / 2, 0)
        self._corner_labels["left"] = (label, position, rotation, kwargs)

    def top_corner_label(self, label, position=None, rotation=0, offset=0.2,
                         **kwargs):
        """
        Sets the label on the bottom axis.

        Parameters
        ----------
        label: String
            The axis label
        position: 3-Tuple of floats, None
            The position of the text label
        rotation: float, 0
            The angle of rotation of the label
        offset: float,
            Used to compute the distance of the label from the axis
        kwargs:
            Any kwargs to pass through to matplotlib.
        """

        if not position:
            position = (-offset / 2, 1 + offset, 0)
        self._corner_labels["top"] = (label, position, rotation, kwargs)

    def annotate(self, text, position, **kwargs):
        ax = self.get_axes()
        p = project_point(position)
        ax.annotate(text, (p[0], p[1]), **kwargs)

    # Boundary and Gridlines

    def boundary(self, scale=None, axes_colors=None, **kwargs):
        # Sometimes you want to draw a bigger boundary
        if not scale:
            scale = self._boundary_scale  # defaults to self._scale
        ax = self.get_axes()
        self.resize_drawing_canvas(scale)
        lines.boundary(scale=scale, ax=ax, axes_colors=axes_colors, **kwargs)

    def gridlines(self, multiple=None, horizontal_kwargs=None, left_kwargs=None,
                  right_kwargs=None, **kwargs):
        ax = self.get_axes()
        scale = self.get_scale()
        lines.gridlines(scale=scale, multiple=multiple,
                        ax=ax, horizontal_kwargs=horizontal_kwargs,
                        left_kwargs=left_kwargs, right_kwargs=right_kwargs,
                        **kwargs)

    # Various Lines

    def line(self, p1, p2, **kwargs):
        ax = self.get_axes()
        lines.line(ax, p1, p2, **kwargs)

    def horizontal_line(self, i, **kwargs):
        ax = self.get_axes()
        scale = self.get_scale()
        lines.horizontal_line(ax, scale, i, **kwargs)

    def left_parallel_line(self, i, **kwargs):
        ax = self.get_axes()
        scale = self.get_scale()
        lines.left_parallel_line(ax, scale, i, **kwargs)

    def right_parallel_line(self, i, **kwargs):
        ax = self.get_axes()
        scale = self.get_scale()
        lines.right_parallel_line(ax, scale, i, **kwargs)

    # Matplotlib passthroughs

    def close(self):
        fig = self.get_figure()
        pyplot.close(fig)

    def legend(self, *args, **kwargs):
        ax = self.get_axes()
        ax.legend(*args, **kwargs)

    def savefig(self, filename, **kwargs):
        self._redraw_labels()
        fig = self.get_figure()
        if 'dpi' not in kwargs:
            kwargs['dpi'] = 200
        fig.savefig(filename, **kwargs)

    def show(self):
        self._redraw_labels()
        pyplot.show()

    # Axis ticks

    def clear_matplotlib_ticks(self, axis="both"):
        """Clears the default matplotlib ticks."""
        ax = self.get_axes()
        plotting.clear_matplotlib_ticks(ax=ax, axis=axis)

    def get_ticks_from_axis_limits(self, multiple=1):
        """
        Taking self._axis_limits and self._boundary_scale get the scaled
        ticks for all three axes and store them in self._ticks under the
        keys 'b' for bottom, 'l' for left and 'r' for right axes.
        """
        for k in ['b', 'l', 'r']:
            self._ticks[k] = numpy.linspace(
                self._axis_limits[k][0],
                self._axis_limits[k][1],
                int(self._boundary_scale / float(multiple) + 1)
            ).tolist()

    def set_custom_ticks(self, locations=None, clockwise=False, multiple=1,
                         axes_colors=None, tick_formats=None, **kwargs):
        """
        Having called get_ticks_from_axis_limits, set the custom ticks on the
        plot.
        """
        for k in ['b', 'l', 'r']:
            self.ticks(ticks=self._ticks[k], locations=locations,
                       axis=k, clockwise=clockwise, multiple=multiple,
                       axes_colors=axes_colors, tick_formats=tick_formats,
                       **kwargs)

    def ticks(self, ticks=None, locations=None, multiple=1, axis='blr',
              clockwise=False, axes_colors=None, tick_formats=None, **kwargs):
        ax = self.get_axes()
        scale = self.get_scale()
        lines.ticks(ax, scale, ticks=ticks, locations=locations,
                    multiple=multiple, clockwise=clockwise, axis=axis,
                    axes_colors=axes_colors, tick_formats=tick_formats,
                    **kwargs)

    # Redrawing and resizing

    def resize_drawing_canvas(self, scale=None):
        ax = self.get_axes()
        if not scale:
            scale = self.get_scale()
        plotting.resize_drawing_canvas(ax, scale=scale)

    def _redraw_labels(self):
        """Redraw axis labels, typically after draw or resize events."""
        ax = self.get_axes()
        # Remove any previous labels
        for mpl_object in self._to_remove:
            mpl_object.remove()
        self._to_remove = []
        # Redraw the labels with the appropriate angles
        label_data = list(self._labels.values())
        label_data.extend(self._corner_labels.values())
        for (label, position, rotation, kwargs) in label_data:
            transform = ax.transAxes
            x, y = project_point(position)
            # Calculate the new angle.
            position = numpy.array([x, y])
            new_rotation = ax.transData.transform_angles(
                numpy.array((rotation,)), position.reshape((1, 2)))[0]
            text = ax.text(x, y, label, rotation=new_rotation,
                           transform=transform, horizontalalignment="center",
                           **kwargs)
            text.set_rotation_mode("anchor")
            self._to_remove.append(text)

    def convert_coordinates(self, points, axisorder='blr'):
        """
        Convert data coordinates to simplex coordinates for plotting
        in the case that axis limits have been applied.
        """
        return convert_coordinates_sequence(points,self._boundary_scale,
                                            self._axis_limits, axisorder)

    # Various Plots

    def scatter(self, points, **kwargs):
        ax = self.get_axes()
        permutation = self._permutation
        plot_ = plotting.scatter(points, ax=ax, permutation=permutation,
                                 **kwargs)
        return plot_

    def plot(self, points, **kwargs):
        ax = self.get_axes()
        permutation = self._permutation
        plotting.plot(points, ax=ax, permutation=permutation,
                      **kwargs)

    def plot_colored_trajectory(self, points, cmap=None, **kwargs):
        ax = self.get_axes()
        permutation = self._permutation
        plotting.plot_colored_trajectory(points, cmap=cmap, ax=ax,
                                         permutation=permutation, **kwargs)

    def heatmap(self, data, scale=None, cmap=None, scientific=False,
                style='triangular', colorbar=True, use_rgba=False,
                vmin=None, vmax=None, cbarlabel=None, cb_kwargs=None):
        permutation = self._permutation
        if not scale:
            scale = self.get_scale()
        if style.lower()[0] == 'd':
            self._boundary_scale = scale + 1
        ax = self.get_axes()
        heatmapping.heatmap(data, scale, cmap=cmap, style=style, ax=ax,
                            scientific=scientific, colorbar=colorbar,
                            permutation=permutation, use_rgba=use_rgba,
                            vmin=vmin, vmax=vmax, cbarlabel=cbarlabel,
                            cb_kwargs=cb_kwargs)

    def heatmapf(self, func, scale=None, cmap=None, boundary=True,
                 style='triangular', colorbar=True, scientific=False,
                 vmin=None, vmax=None, cbarlabel=None, cb_kwargs=None):
        if not scale:
            scale = self.get_scale()
        if style.lower()[0] == 'd':
            self._boundary_scale = scale + 1
        permutation = self._permutation
        ax = self.get_axes()
        heatmapping.heatmapf(func, scale, cmap=cmap, style=style,
                             boundary=boundary, ax=ax, scientific=scientific,
                             colorbar=colorbar, permutation=permutation,
                             vmin=vmin, vmax=vmax, cbarlabel=cbarlabel,
                             cb_kwargs=cb_kwargs)

    def set_background_color(self, color="whitesmoke", zorder=-1000, alpha=0.75):
        self._background_parameters = BackgroundParameters(color=color, alpha=alpha, zorder=zorder)
        self._draw_background()

    def _draw_background(self):
        color, alpha, zorder = self._background_parameters
        scale = self.get_scale()
        ax = self.get_axes()

        # Remove any existing background
        if self._background_triangle:
            print(self._background_triangle)
            self._background_triangle.remove()

        # Draw the background
        self._background_triangle = heatmapping.background_color(ax, color, scale, alpha=alpha, zorder=zorder)[0]
<commit_msg>Remove debugging print statement<commit_after>"""
Wrapper class for all ternary plotting functions.
"""

from collections import namedtuple
from functools import partial

import numpy as np
from matplotlib import pyplot as plt

from . import heatmapping
from . import lines
from . import plotting
from .helpers import project_point, convert_coordinates_sequence

BackgroundParameters = namedtuple('BackgroundParameters', ['color', 'alpha', 'zorder'])

def figure(ax=None, scale=None, permutation=None):
    """
    Wraps a Matplotlib AxesSubplot or generates a new one. Emulates matplotlib's
    > figure, ax = plt.subplots()

    Parameters
    ----------
    ax: AxesSubplot, None
        The matplotlib AxesSubplot to wrap
    scale: float, None
        The scale factor of the ternary plot
    """

    ternary_ax = TernaryAxesSubplot(ax=ax, scale=scale, permutation=permutation)
    return ternary_ax.get_figure(), ternary_ax

def mpl_redraw_callback(event, tax):
    """
    Callback to properly rotate and redraw text labels when the plot is drawn
    or resized.

    Parameters
    ----------
    event: a matplotlib event
        either 'resize_event' or 'draw_event'
    tax: TernaryAxesSubplot
         the TernaryAxesSubplot
    """
    tax._redraw_labels()

class TernaryAxesSubplot(object):
    """
    Wrapper for python-ternary and matplotlib figure. Parameters for member
    functions simply pass through to ternary's functions with the same names.
    This class manages the matplotlib axes, the scale, and the boundary scale
    to ease the use of ternary plotting functions.
    """

    def __init__(self, ax=None, scale=None, permutation=None):
        if not scale:
            scale = 1.0
        if ax:
            self.ax = ax
        else:
            _, self.ax = plt.subplots()
        self.set_scale(scale=scale)
        self._permutation = permutation
        self._boundary_scale = scale
        # Container for the axis labels supplied by the user
        self._labels = dict()
        self._corner_labels = dict()
        self._ticks = dict()
        # Container for the redrawing of labels
        self._to_remove = []
        self._connect_callbacks()
        # Background
        self._background_parameters = None
        # Cache for the background triangle object, so it can be removed and redrawn as needed.
        self._background_triangle = None
        self.set_background_color(color="whitesmoke", zorder=-1000, alpha=0.75)

    def _connect_callbacks(self):
        """Connect resize matplotlib callbacks."""
        figure = self.get_figure()
        callback = partial(mpl_redraw_callback, tax=self)
        event_names = ('resize_event', 'draw_event')
        for event_name in event_names:
            figure.canvas.mpl_connect(event_name, callback)

    def __repr__(self):
        return "TernaryAxesSubplot: %s" % self.ax.__hash__()

    def get_axes(self):
        """Returns the underlying matplotlib AxesSubplot object."""
        return self.ax

    def get_figure(self):
        """Return the underlying matplotlib figure object."""
        ax = self.get_axes()
        return ax.get_figure()

    def set_scale(self, scale=None):
        self._scale = scale
        self.resize_drawing_canvas()

    def get_scale(self):
        return self._scale

    def set_axis_limits(self, axis_limits=None):
        """
        Set min and max data limits for each of the three axes.

        axis_limits = dict
            keys are 'b','l' and 'r' for the three axes
            vals are lists of the min and max values for the axis in
            data units.
        """
        self._axis_limits = axis_limits

    def get_axis_limits(self):
        return self._axis_limits

    # Title and Axis Labels

    def set_title(self, title, **kwargs):
        """Sets the title on the underlying matplotlib AxesSubplot."""
        ax = self.get_axes()
        ax.set_title(title, **kwargs)

    def left_axis_label(self, label, position=None,  rotation=60, offset=0.08,
                        **kwargs):
        """
        Sets the label on the left axis.

        Parameters
        ----------
        label: String
            The axis label
        position: 3-Tuple of floats, None
            The position of the text label
        rotation: float, 60
            The angle of rotation of the label
        offset: float,
            Used to compute the distance of the label from the axis
        kwargs:
            Any kwargs to pass through to matplotlib.
        """

        if not position:
            position = (-offset, 3./5, 2./5)
        self._labels["left"] = (label, position, rotation, kwargs)

    def right_axis_label(self, label, position=None, rotation=-60, offset=0.08,
                         **kwargs):

        """
        Sets the label on the right axis.

        Parameters
        ----------
        label: String
            The axis label
        position: 3-Tuple of floats, None
            The position of the text label
        rotation: float, -60
            The angle of rotation of the label
        offset: float,
            Used to compute the distance of the label from the axis
        kwargs:
            Any kwargs to pass through to matplotlib.
        """

        if not position:
            position = (2. / 5 + offset, 3. / 5, 0)
        self._labels["right"] = (label, position, rotation, kwargs)

    def bottom_axis_label(self, label, position=None, rotation=0, offset=0.02,
                          **kwargs):
        """
        Sets the label on the bottom axis.

        Parameters
        ----------
        label: String
            The axis label
        position: 3-Tuple of floats, None
            The position of the text label
        rotation: float, 0
            The angle of rotation of the label
        offset: float,
            Used to compute the distance of the label from the axis
        kwargs:
            Any kwargs to pass through to matplotlib.
        """

        if not position:
            position = (0.5, -offset / 2., 0.5)
        self._labels["bottom"] = (label, position, rotation, kwargs)

    def right_corner_label(self, label, position=None, rotation=0, offset=0.08,
                           **kwargs):
        """
        Sets the label on the right corner (complements left axis).

        Parameters
        ----------
        label: String
            The axis label
        position: 3-Tuple of floats, None
            The position of the text label
        rotation: float, 0
            The angle of rotation of the label
        offset: float,
            Used to compute the distance of the label from the axis
        kwargs:
            Any kwargs to pass through to matplotlib.
        """

        if not position:
            position = (1, offset / 2, 0)
        self._corner_labels["right"] = (label, position, rotation, kwargs)

    def left_corner_label(self, label, position=None, rotation=0, offset=0.08,
                          **kwargs):
        """
        Sets the label on the left corner (complements right axis.)

        Parameters
        ----------
        label: string
            The axis label
        position: 3-Tuple of floats, None
            The position of the text label
        rotation: float, 0
            The angle of rotation of the label
        offset: float,
            Used to compute the distance of the label from the axis
        kwargs:
            Any kwargs to pass through to matplotlib.
        """

        if not position:
            position = (-offset / 2, offset / 2, 0)
        self._corner_labels["left"] = (label, position, rotation, kwargs)

    def top_corner_label(self, label, position=None, rotation=0, offset=0.2,
                         **kwargs):
        """
        Sets the label on the bottom axis.

        Parameters
        ----------
        label: String
            The axis label
        position: 3-Tuple of floats, None
            The position of the text label
        rotation: float, 0
            The angle of rotation of the label
        offset: float,
            Used to compute the distance of the label from the axis
        kwargs:
            Any kwargs to pass through to matplotlib.
        """

        if not position:
            position = (-offset / 2, 1 + offset, 0)
        self._corner_labels["top"] = (label, position, rotation, kwargs)

    def annotate(self, text, position, **kwargs):
        ax = self.get_axes()
        p = project_point(position)
        ax.annotate(text, (p[0], p[1]), **kwargs)

    # Boundary and Gridlines

    def boundary(self, scale=None, axes_colors=None, **kwargs):
        # Sometimes you want to draw a bigger boundary
        if not scale:
            scale = self._boundary_scale  # defaults to self._scale
        ax = self.get_axes()
        self.resize_drawing_canvas(scale)
        lines.boundary(scale=scale, ax=ax, axes_colors=axes_colors, **kwargs)

    def gridlines(self, multiple=None, horizontal_kwargs=None, left_kwargs=None,
                  right_kwargs=None, **kwargs):
        ax = self.get_axes()
        scale = self.get_scale()
        lines.gridlines(scale=scale, multiple=multiple,
                        ax=ax, horizontal_kwargs=horizontal_kwargs,
                        left_kwargs=left_kwargs, right_kwargs=right_kwargs,
                        **kwargs)

    # Various Lines

    def line(self, p1, p2, **kwargs):
        ax = self.get_axes()
        lines.line(ax, p1, p2, **kwargs)

    def horizontal_line(self, i, **kwargs):
        ax = self.get_axes()
        scale = self.get_scale()
        lines.horizontal_line(ax, scale, i, **kwargs)

    def left_parallel_line(self, i, **kwargs):
        ax = self.get_axes()
        scale = self.get_scale()
        lines.left_parallel_line(ax, scale, i, **kwargs)

    def right_parallel_line(self, i, **kwargs):
        ax = self.get_axes()
        scale = self.get_scale()
        lines.right_parallel_line(ax, scale, i, **kwargs)

    # Matplotlib passthroughs

    def close(self):
        fig = self.get_figure()
        plt.close(fig)

    def legend(self, *args, **kwargs):
        ax = self.get_axes()
        ax.legend(*args, **kwargs)

    def savefig(self, filename, **kwargs):
        self._redraw_labels()
        fig = self.get_figure()
        if 'dpi' not in kwargs:
            kwargs['dpi'] = 200
        fig.savefig(filename, **kwargs)

    def show(self):
        self._redraw_labels()
        plt.show()

    # Axis ticks

    def clear_matplotlib_ticks(self, axis="both"):
        """Clears the default matplotlib ticks."""
        ax = self.get_axes()
        plotting.clear_matplotlib_ticks(ax=ax, axis=axis)

    def get_ticks_from_axis_limits(self, multiple=1):
        """
        Taking self._axis_limits and self._boundary_scale get the scaled
        ticks for all three axes and store them in self._ticks under the
        keys 'b' for bottom, 'l' for left and 'r' for right axes.
        """
        for k in ['b', 'l', 'r']:
            self._ticks[k] = np.linspace(
                self._axis_limits[k][0],
                self._axis_limits[k][1],
                int(self._boundary_scale / float(multiple) + 1)
            ).tolist()

    def set_custom_ticks(self, locations=None, clockwise=False, multiple=1,
                         axes_colors=None, tick_formats=None, **kwargs):
        """
        Having called get_ticks_from_axis_limits, set the custom ticks on the
        plot.
        """
        for k in ['b', 'l', 'r']:
            self.ticks(ticks=self._ticks[k], locations=locations,
                       axis=k, clockwise=clockwise, multiple=multiple,
                       axes_colors=axes_colors, tick_formats=tick_formats,
                       **kwargs)

    def ticks(self, ticks=None, locations=None, multiple=1, axis='blr',
              clockwise=False, axes_colors=None, tick_formats=None, **kwargs):
        ax = self.get_axes()
        scale = self.get_scale()
        lines.ticks(ax, scale, ticks=ticks, locations=locations,
                    multiple=multiple, clockwise=clockwise, axis=axis,
                    axes_colors=axes_colors, tick_formats=tick_formats,
                    **kwargs)

    # Redrawing and resizing

    def resize_drawing_canvas(self, scale=None):
        ax = self.get_axes()
        if not scale:
            scale = self.get_scale()
        plotting.resize_drawing_canvas(ax, scale=scale)

    def _redraw_labels(self):
        """Redraw axis labels, typically after draw or resize events."""
        ax = self.get_axes()
        # Remove any previous labels
        for mpl_object in self._to_remove:
            mpl_object.remove()
        self._to_remove = []
        # Redraw the labels with the appropriate angles
        label_data = list(self._labels.values())
        label_data.extend(self._corner_labels.values())
        for (label, position, rotation, kwargs) in label_data:
            transform = ax.transAxes
            x, y = project_point(position)
            # Calculate the new angle.
            position = np.array([x, y])
            new_rotation = ax.transData.transform_angles(
                np.array((rotation,)), position.reshape((1, 2)))[0]
            text = ax.text(x, y, label, rotation=new_rotation,
                           transform=transform, horizontalalignment="center",
                           **kwargs)
            text.set_rotation_mode("anchor")
            self._to_remove.append(text)

    def convert_coordinates(self, points, axisorder='blr'):
        """
        Convert data coordinates to simplex coordinates for plotting
        in the case that axis limits have been applied.
        """
        return convert_coordinates_sequence(points,self._boundary_scale,
                                            self._axis_limits, axisorder)

    # Various Plots

    def scatter(self, points, **kwargs):
        ax = self.get_axes()
        permutation = self._permutation
        plot_ = plotting.scatter(points, ax=ax, permutation=permutation,
                                 **kwargs)
        return plot_

    def plot(self, points, **kwargs):
        ax = self.get_axes()
        permutation = self._permutation
        plotting.plot(points, ax=ax, permutation=permutation,
                      **kwargs)

    def plot_colored_trajectory(self, points, cmap=None, **kwargs):
        ax = self.get_axes()
        permutation = self._permutation
        plotting.plot_colored_trajectory(points, cmap=cmap, ax=ax,
                                         permutation=permutation, **kwargs)

    def heatmap(self, data, scale=None, cmap=None, scientific=False,
                style='triangular', colorbar=True, use_rgba=False,
                vmin=None, vmax=None, cbarlabel=None, cb_kwargs=None):
        permutation = self._permutation
        if not scale:
            scale = self.get_scale()
        if style.lower()[0] == 'd':
            self._boundary_scale = scale + 1
        ax = self.get_axes()
        heatmapping.heatmap(data, scale, cmap=cmap, style=style, ax=ax,
                            scientific=scientific, colorbar=colorbar,
                            permutation=permutation, use_rgba=use_rgba,
                            vmin=vmin, vmax=vmax, cbarlabel=cbarlabel,
                            cb_kwargs=cb_kwargs)

    def heatmapf(self, func, scale=None, cmap=None, boundary=True,
                 style='triangular', colorbar=True, scientific=False,
                 vmin=None, vmax=None, cbarlabel=None, cb_kwargs=None):
        if not scale:
            scale = self.get_scale()
        if style.lower()[0] == 'd':
            self._boundary_scale = scale + 1
        permutation = self._permutation
        ax = self.get_axes()
        heatmapping.heatmapf(func, scale, cmap=cmap, style=style,
                             boundary=boundary, ax=ax, scientific=scientific,
                             colorbar=colorbar, permutation=permutation,
                             vmin=vmin, vmax=vmax, cbarlabel=cbarlabel,
                             cb_kwargs=cb_kwargs)

    def set_background_color(self, color="whitesmoke", zorder=-1000, alpha=0.75):
        self._background_parameters = BackgroundParameters(color=color, alpha=alpha, zorder=zorder)
        self._draw_background()

    def _draw_background(self):
        color, alpha, zorder = self._background_parameters
        scale = self.get_scale()
        ax = self.get_axes()

        # Remove any existing background
        if self._background_triangle:
            self._background_triangle.remove()

        # Draw the background
        self._background_triangle = heatmapping.background_color(ax, color, scale, alpha=alpha, zorder=zorder)[0]
<|endoftext|>


<commit_before>#
# Copyright 2017 GoPro Inc.
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
#

import os
import os.path as op
import subprocess

import pynodegl as ngl
from PySide2 import QtGui, QtCore

from .com import query_inplace
from .misc import get_backend, get_viewport, get_nodegl_tempdir

class Exporter(QtCore.QThread):

    progressed = QtCore.Signal(int)
    failed = QtCore.Signal()
    export_finished = QtCore.Signal()

    def __init__(self, get_scene_func, filename, w, h, extra_enc_args=None, time=None):
        super().__init__()
        self._get_scene_func = get_scene_func
        self._filename = filename
        self._width = w
        self._height = h
        self._extra_enc_args = extra_enc_args if extra_enc_args is not None else []
        self._time = time
        self._cancelled = False

    def run(self):
        filename, width, height = self._filename, self._width, self._height

        if filename.endswith('gif'):
            palette_filename = op.join(get_nodegl_tempdir(), 'palette.png')
            pass1_args = ['-vf', 'palettegen']
            pass2_args = self._extra_enc_args + ['-i', palette_filename, '-lavfi', 'paletteuse']
            ok = self._export(palette_filename, width, height, pass1_args)
            if not ok:
                return
            ok = self._export(filename, width, height, pass2_args)
        else:
            ok = self._export(filename, width, height, self._extra_enc_args)
        if ok:
            self.export_finished.emit()

    def _export(self, filename, width, height, extra_enc_args=None):
        fd_r, fd_w = os.pipe()

        cfg = self._get_scene_func()
        if not cfg:
            self.failed.emit()
            return False

        fps = cfg['framerate']
        duration = cfg['duration']
        samples = cfg['samples']

        cmd = ['ffmpeg', '-r', '%d/%d' % fps,
               '-nostats', '-nostdin',
               '-f', 'rawvideo',
               '-video_size', '%dx%d' % (width, height),
               '-pixel_format', 'rgba',
               '-i', 'pipe:%d' % fd_r]
        if extra_enc_args:
            cmd += extra_enc_args
        cmd += ['-y', filename]

        reader = subprocess.Popen(cmd, pass_fds=(fd_r,))
        os.close(fd_r)

        capture_buffer = bytearray(width * height * 4)

        # node.gl context
        ctx = ngl.Context()
        ctx.configure(
            platform=ngl.PLATFORM_AUTO,
            backend=get_backend(cfg['backend']),
            offscreen=1,
            width=width,
            height=height,
            viewport=get_viewport(width, height, cfg['aspect_ratio']),
            samples=samples,
            clear_color=cfg['clear_color'],
            capture_buffer=capture_buffer,
        )
        ctx.set_scene_from_string(cfg['scene'])

        if self._time is not None:
            ctx.draw(self._time)
            os.write(fd_w, capture_buffer)
            self.progressed.emit(100)
        else:
            # Draw every frame
            nb_frame = int(duration * fps[0] / fps[1])
            for i in range(nb_frame):
                if self._cancelled:
                    break
                time = i * fps[1] / float(fps[0])
                ctx.draw(time)
                os.write(fd_w, capture_buffer)
                self.progressed.emit(i*100 / nb_frame)
            self.progressed.emit(100)

        os.close(fd_w)
        reader.wait()
        return True

    def cancel(self):
        self._cancelled = True

def test_export():
    import sys

    def _get_scene(**cfg_overrides):
        cfg = {
            'pkg': 'pynodegl_utils.examples',
            'scene': ('misc', 'triangle'),
            'duration': 5,
        }
        cfg.update(cfg_overrides)

        ret = query_inplace(query='scene', **cfg)
        if 'error' in ret:
            print(ret['error'])
            return None
        return ret

    def print_progress(progress):
        sys.stdout.write('\r%d%%' % progress)
        sys.stdout.flush()
        if progress == 100:
            sys.stdout.write('\n')

    if len(sys.argv) != 2:
        print('Usage: %s <outfile>' % sys.argv[0])
        sys.exit(0)

    filename = sys.argv[1]
    app = QtGui.QGuiApplication(sys.argv)

    exporter = Exporter(_get_scene, filename, 320, 240)
    exporter.progressed.connect(print_progress)
    exporter.start()
    exporter.wait()

if __name__ == '__main__':
    test_export()
<commit_msg>utils/export: fix imports when running the script directly<commit_after>#
# Copyright 2017 GoPro Inc.
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
#

import os
import os.path as op
import subprocess

import pynodegl as ngl
from PySide2 import QtGui, QtCore

from pynodegl_utils.com import query_inplace
from pynodegl_utils.misc import get_backend, get_viewport, get_nodegl_tempdir

class Exporter(QtCore.QThread):

    progressed = QtCore.Signal(int)
    failed = QtCore.Signal()
    export_finished = QtCore.Signal()

    def __init__(self, get_scene_func, filename, w, h, extra_enc_args=None, time=None):
        super().__init__()
        self._get_scene_func = get_scene_func
        self._filename = filename
        self._width = w
        self._height = h
        self._extra_enc_args = extra_enc_args if extra_enc_args is not None else []
        self._time = time
        self._cancelled = False

    def run(self):
        filename, width, height = self._filename, self._width, self._height

        if filename.endswith('gif'):
            palette_filename = op.join(get_nodegl_tempdir(), 'palette.png')
            pass1_args = ['-vf', 'palettegen']
            pass2_args = self._extra_enc_args + ['-i', palette_filename, '-lavfi', 'paletteuse']
            ok = self._export(palette_filename, width, height, pass1_args)
            if not ok:
                return
            ok = self._export(filename, width, height, pass2_args)
        else:
            ok = self._export(filename, width, height, self._extra_enc_args)
        if ok:
            self.export_finished.emit()

    def _export(self, filename, width, height, extra_enc_args=None):
        fd_r, fd_w = os.pipe()

        cfg = self._get_scene_func()
        if not cfg:
            self.failed.emit()
            return False

        fps = cfg['framerate']
        duration = cfg['duration']
        samples = cfg['samples']

        cmd = ['ffmpeg', '-r', '%d/%d' % fps,
               '-nostats', '-nostdin',
               '-f', 'rawvideo',
               '-video_size', '%dx%d' % (width, height),
               '-pixel_format', 'rgba',
               '-i', 'pipe:%d' % fd_r]
        if extra_enc_args:
            cmd += extra_enc_args
        cmd += ['-y', filename]

        reader = subprocess.Popen(cmd, pass_fds=(fd_r,))
        os.close(fd_r)

        capture_buffer = bytearray(width * height * 4)

        # node.gl context
        ctx = ngl.Context()
        ctx.configure(
            platform=ngl.PLATFORM_AUTO,
            backend=get_backend(cfg['backend']),
            offscreen=1,
            width=width,
            height=height,
            viewport=get_viewport(width, height, cfg['aspect_ratio']),
            samples=samples,
            clear_color=cfg['clear_color'],
            capture_buffer=capture_buffer,
        )
        ctx.set_scene_from_string(cfg['scene'])

        if self._time is not None:
            ctx.draw(self._time)
            os.write(fd_w, capture_buffer)
            self.progressed.emit(100)
        else:
            # Draw every frame
            nb_frame = int(duration * fps[0] / fps[1])
            for i in range(nb_frame):
                if self._cancelled:
                    break
                time = i * fps[1] / float(fps[0])
                ctx.draw(time)
                os.write(fd_w, capture_buffer)
                self.progressed.emit(i*100 / nb_frame)
            self.progressed.emit(100)

        os.close(fd_w)
        reader.wait()
        return True

    def cancel(self):
        self._cancelled = True

def test_export():
    import sys

    def _get_scene(**cfg_overrides):
        cfg = {
            'pkg': 'pynodegl_utils.examples',
            'scene': ('misc', 'triangle'),
            'duration': 5,
        }
        cfg.update(cfg_overrides)

        ret = query_inplace(query='scene', **cfg)
        if 'error' in ret:
            print(ret['error'])
            return None
        return ret

    def print_progress(progress):
        sys.stdout.write('\r%d%%' % progress)
        sys.stdout.flush()
        if progress == 100:
            sys.stdout.write('\n')

    if len(sys.argv) != 2:
        print('Usage: %s <outfile>' % sys.argv[0])
        sys.exit(0)

    filename = sys.argv[1]
    app = QtGui.QGuiApplication(sys.argv)

    exporter = Exporter(_get_scene, filename, 320, 240)
    exporter.progressed.connect(print_progress)
    exporter.start()
    exporter.wait()

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