text:
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""A decoder that performs beam search."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import math
import numpy as np
import tensorflow as tf

import decoder

__all__ = [
    "BeamSearchDecoderOutput",
    "BeamSearchDecoderState",
    "BeamSearchDecoder",
    "tile_batch",
]

class BeamSearchDecoderState(
    collections.namedtuple("BeamSearchDecoderState",
                           ("cell_state", "log_probs", "finished", "lengths",
                            "accumulated_attention_probs", "pred_ids"))):
  pass

class BeamSearchDecoderOutput(
    collections.namedtuple("BeamSearchDecoderOutput",
                           ("scores", "predicted_ids", "parent_ids"))):
  pass

def _tile_batch(t, multiplier):
  """Core single-tensor implementation of tile_batch."""
  t = tf.convert_to_tensor(t, name="t")
  shape_t = tf.shape(t)
  if t.shape.ndims is None or t.shape.ndims < 1:
    raise ValueError("t must have statically known rank")
  tiling = [1] * (t.shape.ndims + 1)
  tiling[1] = multiplier
  tiled_static_batch_size = (
      t.shape[0].value * multiplier if t.shape[0].value is not None else None)
  tiled = tf.tile(tf.expand_dims(t, 1), tiling)
  tiled = tf.reshape(
      tiled, tf.concat(([shape_t[0] * multiplier], shape_t[1:]), 0))
  tiled.set_shape(
      tf.TensorShape([tiled_static_batch_size]).concatenate(
          t.shape[1:]))
  return tiled

def tile_batch(t, multiplier, name=None):
  """Tile the batch dimension of a (possibly nested structure of) tensor(s) t.

  For each tensor t in a (possibly nested structure) of tensors,
  this function takes a tensor t shaped `[batch_size, s0, s1, ...]` composed of
  minibatch entries `t[0], ..., t[batch_size - 1]` and tiles it to have a shape
  `[batch_size * multiplier, s0, s1, ...]` composed of minibatch entries
  `t[0], t[0], ..., t[1], t[1], ...` where each minibatch entry is repeated
  `multiplier` times.

  Args:
    t: `Tensor` shaped `[batch_size, ...]`.
    multiplier: Python int.
    name: Name scope for any created operations.

  Returns:
    A (possibly nested structure of) `Tensor` shaped
    `[batch_size * multiplier, ...]`.

  Raises:
    ValueError: if tensor(s) `t` do not have a statically known rank or
    the rank is < 1.
  """
  flat_t = tf.contrib.framework.nest.flatten(t)
  with tf.name_scope(name, "tile_batch", flat_t + [multiplier]):
    return tf.contrib.framework.nest.map_structure(
        lambda t_: _tile_batch(t_, multiplier), t)

def _check_maybe(t):
  if t.shape.ndims is None:
    raise ValueError(
        "Expected tensor (%s) to have known rank, but ndims == None." % t)

class BeamSearchDecoder(decoder.Decoder):
  """BeamSearch sampling decoder.

    **NOTE** If you are using the `BeamSearchDecoder` with a cell wrapped in
    `AttentionWrapper`, then you must ensure that:

    - The encoder output has been tiled to `beam_width` via
      `tf.contrib.seq2seq.tile_batch` (NOT `tf.tile`).
    - The `batch_size` argument passed to the `zero_state` method of this
      wrapper is equal to `true_batch_size * beam_width`.
    - The initial state created with `zero_state` above contains a
      `cell_state` value containing properly tiled final state from the
      encoder.

    An example:

    ```
    tiled_encoder_outputs = tf.contrib.seq2seq.tile_batch(
        encoder_outputs, multiplier=beam_width)
    tiled_encoder_final_state = tf.contrib.seq2seq.tile_batch(
        encoder_final_state, multiplier=beam_width)
    tiled_sequence_length = tf.contrib.seq2seq.tile_batch(
        sequence_length, multiplier=beam_width)
    attention_mechanism = MyFavoriteAttentionMechanism(
        num_units=attention_depth,
        memory=tiled_inputs,
        memory_sequence_length=tiled_sequence_length)
    attention_cell = AttentionWrapper(cell, attention_mechanism, ...)
    decoder_initial_state = attention_cell.zero_state(
        dtype, batch_size=true_batch_size * beam_width)
    decoder_initial_state = decoder_initial_state.clone(
        cell_state=tiled_encoder_final_state)
    ```

    Meanwhile, with `AttentionWrapper`, coverage penalty is suggested to use
    when computing scores(https://arxiv.org/pdf/1609.08144.pdf). It encourages
    the translation to cover all inputs.
  """

  def __init__(self,
               cell,
               embedding,
               start_tokens,
               end_token,
               initial_state,
               beam_width,
               output_layer=None,
               max_tgt=None,
               length_penalty_weight=0.0,
               coverage_penalty_weight=0.0,
               reorder_tensor_arrays=True,
               dtype=tf.float32):
    """Initialize the BeamSearchDecoder.

    Args:
      cell: An `RNNCell` instance.
      embedding: A callable that takes a vector tensor of `ids` (argmax ids),
        or the `params` argument for `embedding_lookup`.
      start_tokens: `int32` vector shaped `[batch_size]`, the start tokens.
      end_token: `int32` scalar, the token that marks end of decoding.
      initial_state: A (possibly nested tuple of...) tensors and TensorArrays.
      beam_width:  Python integer, the number of beams.
      output_layer: (Optional) An instance of `tf.variable` to represent the
        weights for the dense layer to apply to the RNN output prior
        to storing the result or sampling.
      max_tgt: maximum prediction length.
      length_penalty_weight: Float weight to penalize length. Disabled with 0.0.
      coverage_penalty_weight: Float weight to penalize the coverage of source
        sentence. Disabled with 0.0.
      reorder_tensor_arrays: If `True`, `TensorArray`s' elements within the cell
        state will be reordered according to the beam search path. If the
        `TensorArray` can be reordered, the stacked form will be returned.
        Otherwise, the `TensorArray` will be returned as is. Set this flag to
        `False` if the cell state contains `TensorArray`s that are not amenable
        to reordering.

    Raises:
      TypeError: if `cell` is not an instance of `RNNCell`,
        or `output_layer` is not an instance of `tf.layers.Layer`.
      ValueError: If `start_tokens` is not a vector or
        `end_token` is not a scalar.
    """
    self._cell = cell
    self._output_layer = output_layer
    self._reorder_tensor_arrays = reorder_tensor_arrays

    if callable(embedding):
      self._embedding_fn = embedding
    else:
      self._embedding_fn = (
          lambda ids: tf.cast(tf.nn.embedding_lookup(embedding, ids), dtype))

    self._start_tokens = tf.convert_to_tensor(
        start_tokens, dtype=tf.int32, name="start_tokens")
    if self._start_tokens.get_shape().ndims != 1:
      raise ValueError("start_tokens must be a vector")
    self._end_token = tf.convert_to_tensor(
        end_token, dtype=tf.int32, name="end_token")
    if self._end_token.get_shape().ndims != 0:
      raise ValueError("end_token must be a scalar")

    self._batch_size = start_tokens.shape[0].value
    self._beam_width = beam_width
    self._max_tgt = max_tgt
    self._length_penalty_weight = length_penalty_weight
    self._coverage_penalty_weight = coverage_penalty_weight
    self._initial_cell_state = tf.contrib.framework.nest.map_structure(
        self._split_batch_beams, initial_state)
    self._start_tokens = tf.tile(
        tf.expand_dims(self._start_tokens, 1), [1, self._beam_width])
    self._start_inputs = self._embedding_fn(self._start_tokens)

    self._finished = tf.one_hot(
        tf.zeros([self._batch_size], dtype=tf.int32),
        depth=self._beam_width,
        on_value=False,
        off_value=True,
        dtype=tf.bool)

  @property
  def batch_size(self):
    return self._batch_size

  @property
  def tracks_own_finished(self):
    """The BeamSearchDecoder shuffles its beams and their finished state.

    For this reason, it conflicts with the `dynamic_decode` function's
    tracking of finished states.  Setting this property to true avoids
    early stopping of decoding due to mismanagement of the finished state
    in `dynamic_decode`.

    Returns:
      `True`.
    """
    return True

  @property
  def output_size(self):
    # Return the cell output and the id
    return BeamSearchDecoderOutput(
        scores=tf.TensorShape([self._beam_width]),
        predicted_ids=tf.TensorShape([self._beam_width]),
        parent_ids=tf.TensorShape([self._beam_width]))

  @property
  def output_dtype(self):
    # Assume the dtype of the cell is the output_size structure
    # containing the input_state's first component's dtype.
    # Return that structure and int32 (the id)
    dtype = tf.contrib.framework.nest.flatten(self._initial_cell_state)[0].dtype
    return BeamSearchDecoderOutput(
        scores=tf.contrib.framework.nest.map_structure(
            lambda _: dtype, self._output_layer.shape[-1]),
        predicted_ids=tf.int32,
        parent_ids=tf.int32)

  def initialize(self, name=None):
    """Initialize the decoder.

    Args:
      name: Name scope for any created operations.

    Returns:
      `(finished, start_inputs, initial_state)`.
    """
    finished, start_inputs = self._finished, self._start_inputs

    dtype = tf.contrib.framework.nest.flatten(self._initial_cell_state)[0].dtype
    log_probs = tf.one_hot(  # shape(batch_sz, beam_sz)
        tf.zeros([self._batch_size], dtype=tf.int32),
        depth=self._beam_width,
        on_value=tf.convert_to_tensor(0.0, dtype=dtype),
        off_value=tf.convert_to_tensor(-np.Inf, dtype=dtype),
        dtype=dtype)
    init_attention_probs = get_attention_probs(
        self._initial_cell_state, self._coverage_penalty_weight)
    if init_attention_probs is None:
      init_attention_probs = ()
    init_pred_ids = tf.fill([self._batch_size, self._beam_width, self._max_tgt], self._end_token)

    initial_state = BeamSearchDecoderState(
        cell_state=self._initial_cell_state,
        log_probs=log_probs,
        finished=finished,
        lengths=tf.zeros(
            [self._batch_size, self._beam_width], dtype=tf.int64),
        accumulated_attention_probs=init_attention_probs,
        pred_ids=init_pred_ids)

    return (finished, start_inputs, initial_state)

  def _merge_batch_beams(self, t, s=None):
    """Merges the tensor from a batch of beams into a batch by beams.

    More exactly, t is a tensor of dimension [batch_size, beam_width, s]. We
    reshape this into [batch_size*beam_width, s]

    Args:
      t: Tensor of dimension [batch_size, beam_width, s]
      s: (Possibly known) depth shape.

    Returns:
      A reshaped version of t with dimension [batch_size * beam_width, s].
    """
    if t.shape.ndims < 1:
      return t
    if isinstance(s, tf.Tensor):
      s = tf.contrib.util.constant_value(s)
      if isinstance(s, tf.TensorShape):
        return s
      else:
        s = tf.TensorShape(s)
    else:
      s = tf.TensorShape(s)
    t_shape = tf.shape(t)
    batch_size_beam_width = self._batch_size * self._beam_width
    reshaped_t = tf.reshape(
        t,
        tf.concat(([self._batch_size * self._beam_width], t_shape[2:]), 0))
    reshaped_t.set_shape(
        (tf.TensorShape([batch_size_beam_width]).concatenate(s)))
    return reshaped_t

  def _split_batch_beams(self, t, s=None):
    """Splits the tensor from a batch by beams into a batch of beams.

    More exactly, t is a tensor of dimension [batch_size*beam_width, s]. We
    reshape this into [batch_size, beam_width, s]

    Args:
      t: Tensor of dimension [batch_size*beam_width, s].
      s: (Possibly known) depth shape.

    Returns:
      A reshaped version of t with dimension [batch_size, beam_width, s].

    Raises:
      ValueError: If, after reshaping, the new tensor is not shaped
        `[batch_size, beam_width, s]` (assuming batch_size and beam_width
        are known statically).
    """
    if t.shape.ndims < 1:
      return t
    if isinstance(s, tf.Tensor):
      s = tf.TensorShape(tf.contrib.util.constant_value(s))
    else:
      s = tf.TensorShape(s)
    t_shape = tf.shape(t)
    reshaped_t = tf.reshape(
        t,
        tf.concat(([self._batch_size, self._beam_width], t_shape[1:]), 0))
    expected_reshaped_shape = tf.TensorShape(
        [self._batch_size, self._beam_width]).concatenate(s)
    if not reshaped_t.shape.is_compatible_with(expected_reshaped_shape):
      raise ValueError("Unexpected behavior when reshaping between beam width "
                       "and batch size.  The reshaped tensor has shape: %s.  "
                       "We expected it to have shape "
                       "(batch_size, beam_width, depth) == %s.  Perhaps you "
                       "forgot to create a zero_state with "
                       "batch_size=encoder_batch_size * beam_width?" %
                       (reshaped_t.shape, expected_reshaped_shape))
    reshaped_t.set_shape(expected_reshaped_shape)
    return reshaped_t

  def step(self, time, inputs, state, name=None):
    """Perform a decoding step.

    Args:
      time: scalar `int32` tensor.
      inputs: A (structure of) input tensors.
      state: A (structure of) state tensors and TensorArrays.
      name: Name scope for any created operations.

    Returns:
      `(outputs, next_state, next_inputs, finished)`.
    """
    batch_size = self._batch_size
    beam_width = self._beam_width
    end_token = self._end_token
    length_penalty_weight = self._length_penalty_weight
    coverage_penalty_weight = self._coverage_penalty_weight

    with tf.name_scope(name, "BeamSearchDecoderStep", (time, inputs, state)):
      cell_state = state.cell_state
      inputs = tf.contrib.framework.nest.map_structure(
          lambda inp: self._merge_batch_beams(inp, s=inp.shape[2:]), inputs)
      cell_state = tf.contrib.framework.nest.map_structure(
          self._merge_batch_beams, cell_state)
      next_cell_state, cell_outputs = self._cell(inputs, cell_state)
      cell_outputs = tf.contrib.framework.nest.map_structure(
          lambda out: self._split_batch_beams(out, out.shape[1:]), cell_outputs)
      next_cell_state = tf.contrib.framework.nest.map_structure(
          self._split_batch_beams, next_cell_state)

      if self._output_layer is not None:
        output_shape = cell_outputs.shape.as_list()[:-1] + [
            self._output_layer.shape[-1]
        ]
        cell_outputs = tf.matmul(
            tf.reshape(cell_outputs, [-1, cell_outputs.shape.as_list()[-1]]),
            tf.cast(self._output_layer, cell_outputs.dtype))
        cell_outputs = tf.reshape(cell_outputs, output_shape)

      beam_search_output, beam_search_state = _beam_search_step(
          time=time,
          logits=cell_outputs,
          next_cell_state=next_cell_state,
          beam_state=state,
          batch_size=batch_size,
          beam_width=beam_width,
          end_token=end_token,
          length_penalty_weight=length_penalty_weight,
          coverage_penalty_weight=coverage_penalty_weight,
          max_tgt=self._max_tgt)

      finished = beam_search_state.finished
      sample_ids = beam_search_output.predicted_ids
      next_inputs = tf.cond(
          tf.reduce_all(finished), lambda: self._start_inputs,
          lambda: self._embedding_fn(sample_ids))

    return (beam_search_output, beam_search_state, next_inputs, finished)

def create_make_unique(inputs):
  if inputs.shape.ndims != 2:
    raise ValueError("Input of top_k_with_unique must be rank-2 "
                     "but got: %s" % inputs.shape)
  height = inputs.shape[0]
  width = inputs.shape[1]
  zeros = tf.zeros([height, width], dtype=tf.int32)

  # count_mask is used to mask away the low order bits to ensure that every
  # element is distinct.
  log2_ceiling = int(math.ceil(math.log(float(int(width)), 2)))
  next_power_of_two = 1 << log2_ceiling
  count_mask = ~(next_power_of_two - 1)
  count_mask_r0 = tf.constant(count_mask)
  count_mask_r2 = tf.fill([height, width], count_mask_r0)

  # smallest_normal is the bit representation of the smallest
  # positive normal floating point number. The sign is zero,
  # exponent is one, and the fraction is zero.
  smallest_normal = 1 << 23
  smallest_normal_r0 = tf.constant(smallest_normal, dtype=tf.int32)
  smallest_normal_r2 = tf.fill([height, width], smallest_normal_r0)

  # Used to mask away the sign bit when computing the absolute value.
  low_bit_mask = ~(1 << 31)
  low_bit_mask_r0 = tf.constant(low_bit_mask, dtype=tf.int32)
  low_bit_mask_r2 = tf.fill([height, width], low_bit_mask_r0)

  iota = tf.tile(tf.expand_dims(tf.range(width, dtype=tf.int32), 0),
                 [height, 1])

  # Compare the absolute value with positive zero to handle negative zero.
  #
  # Pseudocode: input_no_zeros = abs(input) == 0 ? FLT_MIN : input
  input_r2 = tf.bitcast(inputs, tf.int32)
  abs_r2 = tf.bitwise.bitwise_and(input_r2, low_bit_mask_r2)
  if_zero_r2 = tf.equal(abs_r2, zeros)
  smallest_normal_preserving_sign_r2 = tf.bitwise.bitwise_or(
      input_r2, smallest_normal_r2)
  input_no_zeros_r2 = tf.where(
      if_zero_r2, smallest_normal_preserving_sign_r2, input_r2)

  # Discard the low-order bits and replace with iota.
  and_r2 = tf.bitwise.bitwise_and(input_no_zeros_r2, count_mask_r2)
  or_r2 = tf.bitwise.bitwise_or(and_r2, iota)
  return tf.bitcast(or_r2, tf.float32)

def create_topk_unique(inputs, k):
  height = inputs.shape[0]
  width = inputs.shape[1]
  neg_inf_r0 = tf.constant(-np.inf, dtype=tf.float32)
  ones = tf.ones([height, width], dtype=tf.float32)
  neg_inf_r2 = ones * neg_inf_r0
  inputs = tf.where(tf.is_nan(inputs), neg_inf_r2, inputs)

  tmp = inputs
  topk_r2 = tf.zeros([height, k], dtype=tf.float32)
  for i in range(k):
    kth_order_statistic = tf.reduce_max(tmp, axis=1, keepdims=True)
    k_mask = tf.tile(tf.expand_dims(tf.equal(tf.range(k), tf.fill([k], i)), 0),
                     [height, 1])
    topk_r2 = tf.where(k_mask, tf.tile(kth_order_statistic, [1, k]), topk_r2)
    ge_r2 = tf.greater_equal(inputs, tf.tile(kth_order_statistic, [1, width]))
    tmp = tf.where(ge_r2, neg_inf_r2, inputs)

  log2_ceiling = int(math.ceil(math.log(float(int(width)), 2)))
  next_power_of_two = 1 << log2_ceiling
  count_mask = next_power_of_two - 1
  mask_r0 = tf.constant(count_mask)
  mask_r2 = tf.fill([height, k], mask_r0)
  topk_r2_s32 = tf.bitcast(topk_r2, tf.int32)
  topk_indices_r2 = tf.bitwise.bitwise_and(topk_r2_s32, mask_r2)
  return topk_r2, topk_indices_r2

def top_k_with_unique(inputs, k):
  unique_inputs = create_make_unique(tf.cast(inputs, tf.float32))
  top_values, indices = create_topk_unique(unique_inputs, k)
  top_values = tf.cast(top_values, inputs.dtype)
  return top_values, indices

def _beam_search_step(time, logits, next_cell_state, beam_state, batch_size,
                      beam_width, end_token, length_penalty_weight,
                      coverage_penalty_weight, max_tgt):
  """Performs a single step of Beam Search Decoding.

  Args:
    time: Beam search time step, should start at 0. At time 0 we assume
      that all beams are equal and consider only the first beam for
      continuations.
    logits: Logits at the current time step. A tensor of shape
      `[batch_size, beam_width, vocab_size]`
    next_cell_state: The next state from the cell, e.g. an instance of
      AttentionWrapperState if the cell is attentional.
    beam_state: Current state of the beam search.
      An instance of `BeamSearchDecoderState`.
    batch_size: The batch size for this input.
    beam_width: Python int.  The size of the beams.
    end_token: The int32 end token.
    length_penalty_weight: Float weight to penalize length. Disabled with 0.0.
    coverage_penalty_weight: Float weight to penalize the coverage of source
      sentence. Disabled with 0.0.
    max_tgt: maximum prediction length.

  Returns:
    A new beam state.
  """

  # Calculate the current lengths of the predictions
  prediction_lengths = beam_state.lengths
  previously_finished = beam_state.finished
  not_finished = tf.logical_not(previously_finished)

  # Calculate the total log probs for the new hypotheses
  # Final Shape: [batch_size, beam_width, vocab_size]
  step_log_probs = tf.nn.log_softmax(logits)
  step_log_probs = _mask_probs(step_log_probs, end_token, previously_finished)
  total_probs = tf.expand_dims(beam_state.log_probs, 2) + step_log_probs

  # Calculate the continuation lengths by adding to all continuing beams.
  vocab_size = logits.shape[-1].value or tf.shape(logits)[-1]
  lengths_to_add = tf.one_hot(
      indices=tf.fill([batch_size, beam_width], end_token),
      depth=vocab_size,
      on_value=np.int64(0),
      off_value=np.int64(1),
      dtype=tf.int64)
  add_mask = tf.to_int64(not_finished)
  lengths_to_add *= tf.expand_dims(add_mask, 2)
  new_prediction_lengths = (
      lengths_to_add + tf.expand_dims(prediction_lengths, 2))

  # Calculate the accumulated attention probabilities if coverage penalty is
  # enabled.
  accumulated_attention_probs = None
  attention_probs = get_attention_probs(
      next_cell_state, coverage_penalty_weight)
  if attention_probs is not None:
    attention_probs *= tf.expand_dims(
        tf.cast(not_finished, attention_probs.dtype), 2)
    accumulated_attention_probs = (
        beam_state.accumulated_attention_probs + attention_probs)

  batch_finished = tf.reduce_all(
      previously_finished, axis=1, keepdims=True)
  any_batch_finished = tf.reduce_any(batch_finished)
  batch_finished = tf.tile(tf.expand_dims(batch_finished, 2),
                           [1, beam_width, vocab_size])
  def _normalized_scores():
    return _get_scores(
        log_probs=total_probs,
        sequence_lengths=new_prediction_lengths,
        length_penalty_weight=length_penalty_weight,
        coverage_penalty_weight=coverage_penalty_weight,
        finished=batch_finished,
        accumulated_attention_probs=accumulated_attention_probs)

  # Normalize the scores of finished batches.
  scores = tf.cond(any_batch_finished, _normalized_scores, lambda: total_probs)

  time = tf.convert_to_tensor(time, name="time")
  # During the first time step we only consider the initial beam
  scores_flat = tf.reshape(scores, [batch_size, -1])

  # Pick the next beams according to the specified successors function
  next_beam_scores, word_indices = top_k_with_unique(scores_flat, beam_width)

  next_beam_scores.set_shape([batch_size, beam_width])
  word_indices.set_shape([batch_size, beam_width])

  # Pick out the probs, beam_ids, and states according to the chosen predictions
  batch_ids = tf.expand_dims(
      tf.tile(tf.expand_dims(tf.range(batch_size), 1), [1, beam_width]), 2)
  indices = tf.concat([batch_ids, tf.expand_dims(word_indices, 2)], -1)
  next_beam_probs = tf.gather_nd(
      tf.reshape(total_probs, [batch_size, -1]), indices)

  # Note: just doing the following
  #   tf.to_int32(word_indices % vocab_size,
  #       name="next_beam_word_ids")
  # would be a lot cleaner but for reasons unclear, that hides the results of
  # the op which prevents capturing it with tfdbg debug ops.
  raw_next_word_ids = tf.mod(
      word_indices, vocab_size, name="next_beam_word_ids")
  next_word_ids = tf.to_int32(raw_next_word_ids)
  next_beam_ids = tf.div(word_indices, vocab_size, name="next_beam_parent_ids")

  # Append new ids to current predictions
  previously_finished = _tensor_gather_helper(
      gather_indices=next_beam_ids,
      gather_from=previously_finished,
      batch_size=batch_size,
      range_size=beam_width,
      gather_shape=[-1])
  next_finished = tf.logical_or(
      previously_finished,
      tf.equal(next_word_ids, end_token),
      name="next_beam_finished")

  # Calculate the length of the next predictions.
  # 1. Finished beams remain unchanged.
  # 2. Beams that are now finished (EOS predicted) have their length
  #    increased by 1.
  # 3. Beams that are not yet finished have their length increased by 1.
  lengths_to_add = tf.to_int64(tf.logical_not(previously_finished))
  next_prediction_len = _tensor_gather_helper(
      gather_indices=next_beam_ids,
      gather_from=beam_state.lengths,
      batch_size=batch_size,
      range_size=beam_width,
      gather_shape=[-1])
  next_prediction_len += lengths_to_add
  next_accumulated_attention_probs = ()
  if accumulated_attention_probs is not None:
    next_accumulated_attention_probs = _tensor_gather_helper(
        gather_indices=next_beam_ids,
        gather_from=accumulated_attention_probs,
        batch_size=batch_size,
        range_size=beam_width,
        gather_shape=[batch_size * beam_width, -1],
        name="next_accumulated_attention_probs")
  next_pred_ids = _tensor_gather_helper(
      gather_indices=next_beam_ids,
      gather_from=beam_state.pred_ids,
      batch_size=batch_size,
      range_size=beam_width,
      gather_shape=[batch_size * beam_width, -1],
      name="pred_ids")

  # Add next_word_ids to next_pred_ids.
  next_pred_ids = tf.transpose(next_pred_ids, [2, 0, 1])
  cur_time = tf.tile(tf.reshape(time, [1]), [max_tgt])
  time_mask = tf.equal(tf.range(max_tgt), cur_time)
  time_mask = tf.tile(tf.reshape(
      time_mask, [max_tgt, 1, 1]), [1, batch_size, beam_width])
  cur_time_ids = tf.tile(tf.reshape(next_word_ids, [1, batch_size, beam_width]),
                         [max_tgt, 1, 1])
  next_pred_ids = tf.where(time_mask, cur_time_ids, next_pred_ids)
  next_pred_ids = tf.transpose(next_pred_ids, [1, 2, 0])

  # Pick out the cell_states according to the next_beam_ids. We use a
  # different gather_shape here because the cell_state tensors, i.e.
  # the tensors that would be gathered from, all have dimension
  # greater than two and we need to preserve those dimensions.
  # pylint: disable=g-long-lambda
  next_cell_state = tf.contrib.framework.nest.map_structure(
      lambda gather_from: _maybe_tensor_gather_helper(
          gather_indices=next_beam_ids,
          gather_from=gather_from,
          batch_size=batch_size,
          range_size=beam_width,
          gather_shape=[batch_size * beam_width, -1]),
      next_cell_state)
  # pylint: enable=g-long-lambda

  next_state = BeamSearchDecoderState(
      cell_state=next_cell_state,
      log_probs=next_beam_probs,
      lengths=next_prediction_len,
      finished=next_finished,
      accumulated_attention_probs=next_accumulated_attention_probs,
      pred_ids=next_pred_ids)

  output = BeamSearchDecoderOutput(
      scores=next_beam_scores,
      predicted_ids=next_word_ids,
      parent_ids=next_beam_ids)

  return output, next_state

def get_attention_probs(next_cell_state, coverage_penalty_weight):
  """Get attention probabilities from the cell state.

  Args:
    next_cell_state: The next state from the cell, e.g. an instance of
      AttentionWrapperState if the cell is attentional.
    coverage_penalty_weight: Float weight to penalize the coverage of source
      sentence. Disabled with 0.0.

  Returns:
    The attention probabilities with shape `[batch_size, beam_width, max_time]`
    if coverage penalty is enabled. Otherwise, returns None.

  Raises:
    ValueError: If no cell is attentional but coverage penalty is enabled.
  """
  if coverage_penalty_weight == 0.0:
    return None

  # Attention probabilities of each attention layer. Each with shape
  # `[batch_size, beam_width, max_time]`.
  probs_per_attn_layer = []
  for state in next_cell_state:
    if "attention" in state:
      probs_per_attn_layer.append(attention_probs_from_attn_state(state))

  if not probs_per_attn_layer:
    raise ValueError(
        "coverage_penalty_weight must be 0.0 if no cell is attentional.")

  if len(probs_per_attn_layer) == 1:
    attention_probs = probs_per_attn_layer[0]
  else:
    # Calculate the average attention probabilities from all attention layers.
    attention_probs = [
        tf.expand_dims(prob, -1) for prob in probs_per_attn_layer]
    attention_probs = tf.concat(attention_probs, -1)
    attention_probs = tf.reduce_mean(attention_probs, -1)

  return attention_probs

def _get_scores(log_probs, sequence_lengths, length_penalty_weight,
                coverage_penalty_weight, finished, accumulated_attention_probs):
  """Calculates scores for beam search hypotheses.

  Args:
    log_probs: The log probabilities with shape
      `[batch_size, beam_width, vocab_size]`.
    sequence_lengths: The array of sequence lengths.
    length_penalty_weight: Float weight to penalize length. Disabled with 0.0.
    coverage_penalty_weight: Float weight to penalize the coverage of source
      sentence. Disabled with 0.0.
    finished: A boolean tensor of shape `[batch_size, beam_width, vocab_size]`
      that specifies which elements in the beam are finished already.
    accumulated_attention_probs: Accumulated attention probabilities up to the
      current time step, with shape `[batch_size, beam_width, max_time]` if
      coverage_penalty_weight is not 0.0.

  Returns:
    The scores normalized by the length_penalty and coverage_penalty.

  Raises:
    ValueError: accumulated_attention_probs is None when coverage penalty is
      enabled.
  """
  length_penalty_ = _length_penalty(
      sequence_lengths=sequence_lengths, penalty_factor=length_penalty_weight,
      dtype=log_probs.dtype)

  if coverage_penalty_weight == 0.0:
    return tf.where(
      finished, log_probs / length_penalty_, log_probs)

  coverage_penalty_weight = tf.convert_to_tensor(
      coverage_penalty_weight, name="coverage_penalty_weight",
      dtype=log_probs.dtype)
  if coverage_penalty_weight.shape.ndims != 0:
    raise ValueError("coverage_penalty_weight should be a scalar, "
                     "but saw shape: %s" % coverage_penalty_weight.shape)

  if accumulated_attention_probs is None:
    raise ValueError(
        "accumulated_attention_probs can be None only if coverage penalty is "
        "disabled.")

  # Add source sequence length mask before computing coverage penalty.
  accumulated_attention_probs = tf.where(
      tf.equal(accumulated_attention_probs, 0.0),
      tf.ones_like(accumulated_attention_probs),
      accumulated_attention_probs)

  # coverage penalty =
  #     sum over `max_time` {log(min(accumulated_attention_probs, 1.0))}
  coverage_penalty = tf.reduce_sum(
      tf.log(tf.minimum(accumulated_attention_probs, 1.0)), 2)
  # Apply coverage penalty to finished predictions.
  weighted_coverage_penalty = coverage_penalty * coverage_penalty_weight
  # Reshape from [batch_size, beam_width] to [batch_size, beam_width, 1]
  weighted_coverage_penalty = tf.expand_dims(
      weighted_coverage_penalty, 2)

  # Normalize the scores of finished predictions.
  return tf.where(
      finished, log_probs / length_penalty_ + weighted_coverage_penalty,
      log_probs)

def attention_probs_from_attn_state(attention_state):
  """Calculates the average attention probabilities.

  Args:
    attention_state: An instance of `AttentionWrapperState`.

  Returns:
    The attention probabilities in the given AttentionWrapperState.
    If there're multiple attention mechanisms, return the average value from
    all attention mechanisms.
  """
  # Attention probabilities over time steps, with shape
  # `[batch_size, beam_width, max_time]`.
  attention_probs = attention_state["alignments"]
  if isinstance(attention_probs, tuple):
    attention_probs = [
        tf.expand_dims(prob, -1) for prob in attention_probs]
    attention_probs = tf.concat(attention_probs, -1)
    attention_probs = tf.reduce_mean(attention_probs, -1)
  return attention_probs

def _length_penalty(sequence_lengths, penalty_factor, dtype):
  """Calculates the length penalty. See https://arxiv.org/abs/1609.08144.

  Returns the length penalty tensor:
  ```
  [(5+sequence_lengths)/6]**penalty_factor
  ```
  where all operations are performed element-wise.

  Args:
    sequence_lengths: `Tensor`, the sequence lengths of each hypotheses.
    penalty_factor: A scalar that weights the length penalty.
    dtype: dtype of result.

  Returns:
    If the penalty is `0`, returns the scalar `1.0`.  Otherwise returns
    the length penalty factor, a tensor with the same shape as
    `sequence_lengths`.
  """
  penalty_factor = tf.convert_to_tensor(
      penalty_factor, name="penalty_factor", dtype=dtype)
  penalty_factor.set_shape(())  # penalty should be a scalar.
  static_penalty = tf.contrib.util.constant_value(penalty_factor)
  if static_penalty is not None and static_penalty == 0:
    return 1.0
  length_penalty_const = 5.0
  return tf.div((length_penalty_const + tf.cast(sequence_lengths, dtype))
                **penalty_factor, (length_penalty_const + 1.)**penalty_factor)

def _mask_probs(probs, eos_token, finished):
  """Masks log probabilities.

  The result is that finished beams allocate all probability mass to eos and
  unfinished beams remain unchanged.

  Args:
    probs: Log probabilities of shape `[batch_size, beam_width, vocab_size]`
    eos_token: An int32 id corresponding to the EOS token to allocate
      probability to.
    finished: A boolean tensor of shape `[batch_size, beam_width]` that
      specifies which elements in the beam are finished already.

  Returns:
    A tensor of shape `[batch_size, beam_width, vocab_size]`, where unfinished
    beams stay unchanged and finished beams are replaced with a tensor with all
    probability on the EOS token.
  """
  vocab_size = tf.shape(probs)[2]
  # All finished examples are replaced with a vector that has all
  # probability on EOS
  finished_row = tf.one_hot(
      eos_token,
      vocab_size,
      dtype=probs.dtype,
      on_value=tf.convert_to_tensor(0., dtype=probs.dtype),
      off_value=probs.dtype.min)
  finished_probs = tf.tile(
      tf.reshape(finished_row, [1, 1, -1]),
      tf.concat([tf.shape(finished), [1]], 0))
  finished_mask = tf.tile(
      tf.expand_dims(finished, 2), [1, 1, vocab_size])

  return tf.where(finished_mask, finished_probs, probs)

def _maybe_tensor_gather_helper(gather_indices, gather_from, batch_size,
                                range_size, gather_shape):
  """Maybe applies _tensor_gather_helper.

  This applies _tensor_gather_helper when the gather_from dims is at least as
  big as the length of gather_shape. This is used in conjunction with nest so
  that we don't apply _tensor_gather_helper to inapplicable values like scalars.

  Args:
    gather_indices: The tensor indices that we use to gather.
    gather_from: The tensor that we are gathering from.
    batch_size: The batch size.
    range_size: The number of values in each range. Likely equal to beam_width.
    gather_shape: What we should reshape gather_from to in order to preserve the
      correct values. An example is when gather_from is the attention from an
      AttentionWrapperState with shape [batch_size, beam_width, attention_size].
      There, we want to preserve the attention_size elements, so gather_shape is
      [batch_size * beam_width, -1]. Then, upon reshape, we still have the
      attention_size as desired.

  Returns:
    output: Gathered tensor of shape tf.shape(gather_from)[:1+len(gather_shape)]
      or the original tensor if its dimensions are too small.
  """
  if isinstance(gather_from, tf.TensorArray):
    return gather_from
  _check_maybe(gather_from)
  if gather_from.shape.ndims >= len(gather_shape):
    return _tensor_gather_helper(
        gather_indices=gather_indices,
        gather_from=gather_from,
        batch_size=batch_size,
        range_size=range_size,
        gather_shape=gather_shape)
  else:
    return gather_from

def _tensor_gather_helper(gather_indices,
                          gather_from,
                          batch_size,
                          range_size,
                          gather_shape,
                          name=None):
  """Helper for gathering the right indices from the tensor.

  This works by reshaping gather_from to gather_shape (e.g. [-1]) and then
  gathering from that according to the gather_indices, which are offset by
  the right amounts in order to preserve the batch order.

  Args:
    gather_indices: The tensor indices that we use to gather.
    gather_from: The tensor that we are gathering from.
    batch_size: The input batch size.
    range_size: The number of values in each range. Likely equal to beam_width.
    gather_shape: What we should reshape gather_from to in order to preserve the
      correct values. An example is when gather_from is the attention from an
      AttentionWrapperState with shape [batch_size, beam_width, attention_size].
      There, we want to preserve the attention_size elements, so gather_shape is
      [batch_size * beam_width, -1]. Then, upon reshape, we still have the
      attention_size as desired.
    name: The tensor name for set of operations. By default this is
      'tensor_gather_helper'. The final output is named 'output'.

  Returns:
    output: Gathered tensor of shape tf.shape(gather_from)[:1+len(gather_shape)]
  """
  with tf.name_scope(name, "tensor_gather_helper"):
    range_ = tf.expand_dims(tf.range(batch_size) * range_size, 1)
    gather_indices = tf.reshape(gather_indices + range_, [-1])
    output = tf.gather(
        tf.reshape(gather_from, gather_shape), gather_indices)
    final_shape = tf.shape(gather_from)[:1 + len(gather_shape)]
    final_static_shape = (
        tf.TensorShape([batch_size]).concatenate(
            gather_from.shape[1:1 + len(gather_shape)]))
    output = tf.reshape(output, final_shape, name="output")
    output.set_shape(final_static_shape)
    return output


text:
import numpy as np
from tkinter import filedialog
import sys
import os
import matplotlib.pyplot as plt
import pyproj
from pydelft import grd

#--------------------------------------------------------------------------------------------------
# READ DELFT DEPTH FILE
class dep():
    """Bathymetry file.  See A.3.4 in Delft3D-FLOW user manual.

    Attributes:
        depth       array of depths, formatted to match grid mxn
        grid        class, from the grid file associated with the depth file
        delimiter   the delimiter type (default is tab )
        filename    the full path of the .grd file
    """
    def __init__(self, fname = None, grid_fname = None):
        depth = ''
        grid = ''
        filename = ''
        delimiter = ''

    def read_dep(self, fname = None, grd_fname = None):
        if not fname:
            fname = filedialog.askopenfilename()
        else:
            fname = fname

        # Read lines
        f = open(fname, 'r')
        lines = f.readlines()

        if '\t' in lines[0]:
            delimiter = '\t'
        else:
            delimiter = 'space'

        z = []

        if delimiter != 'space':
            for line in lines:
                for s in line.split('%s' % delimiter):
                    z.append(float(s))
        else:
            for line in lines:
                for s in line.split():
                    z.append(float(s))

        # Reshape depth
        if grd_fname:
            grid = grd.grd()
            grid.read_grd(grd_fname)
            z = np.ma.masked_equal(z, -999.)
            z = np.ma.compressed(z)
            z = np.reshape(z, (grid.m, grid.n))
        else:
            if delimiter != 'space':
                z = np.reshape(z, (np.size(lines),
                    np.size(line.split('%s' % delimiter))))
                z = z[:-1, :-1] # get rid of -999. Nans
                z = np.reshape(z, (np.size(lines)-1, np.size(line.split())-1))

                z = z[:-1, :-1] # get rid of -999. Nans
            else:
                a = np.where(np.array(z) == -999.)
                cols = np.max(np.ediff1d(a))
                rows = np.size(np.array_split(a[0],np.where(np.diff(a[0])!=1)[0]+1))+1
                z = np.reshape(z, (rows, cols))
                z = z[:-1, :-1]

        # update the class
        self.depth = z
        self.filename = fname
        self.delimiter = delimiter

        fname = None

        f.close()

    def write_dep(self, Z, fname = None, grid_fname = None):
        if not fname:
            fname = filedialog.asksaveasfile()
        else:
            fname = fname

        if not grid_fname:
            if np.ndim(Z) == 1:
                m = np.shape(Z)[0]
                n = 1
            else:
                m = np.shape(Z)[0]
                n = np.shape(Z)[1]
        else:
            grid_fname = grid_fname
            grid = grd.grd()
            grid.read_grd(fname = grid_fname)
            m = grid.m
            n = grid.n

        Z = np.array([np.append(i,-999.) for i in Z])
        Z = np.insert(Z, -1, np.ones(np.shape(Z[0]))*-999.)
        Z = Z.reshape((m+1, n+1))

        np.savetxt(fname, Z, delimiter = '\t', fmt = '%.3f')
        print('saved depth file: %s' %os.path.basename(fname))

        self.read_dep(fname)

    def plot_dep(self):
        plt.pcolormesh(self.grid.x, self.grid.y, self.depth,
                        edgecolor = 'none', linewidth = 0.0005)
        plt.axis('equal')
        plt.show()

ext: py

max_stars_repo_name: spmls/pydelft

max_issues_repo_name: spmls/pydelft

max_forks_repo_name: spmls/pydelft

reason: import numpy


text:
from typing import List, Any
import numpy as np
from scipy.sparse import *
from sklearn.ensemble import RandomForestClassifier

class Scoring:
    def __init__(self, communities: List[Any], data: np.ndarray):
        self.communities = communities
        self.data = data

    def get_scores(self):
        scores = []
        objs, n_features = self.data.shape

        clf = RandomForestClassifier(random_state=0)
        clustering = np.zeros(objs)

        for _, values in enumerate(self.communities):
            for idx in values:
                clustering[idx] = 1

            clf.fit(self.data, clustering)
            scores.append(clf.feature_importances_)
            # scores.append(clf.coef_[0])

            clustering = np.zeros(objs)

        # return np.mean(scores, axis=0)
        return np.linalg.norm(scores, axis=0)

reason: import numpy,from scipy


text:
import numpy as np
import sympy as sp
from sympy.core import S, Dummy, pi
from sympy.functions.combinatorial.factorials import factorial
from sympy.functions.special.gamma_functions import gamma
from sympy.functions.elementary.miscellaneous import sqrt
from sympy.polys.orthopolys import (legendre_poly, laguerre_poly,
                                    hermite_poly, jacobi_poly)
from sympy.polys.rootoftools import RootOf

def gauss_hermite(n, n_digits):
    r"""
    Computes the Gauss-Hermite quadrature [1]_ points and weights.

    The Gauss-Hermite quadrature approximates the integral:

    .. math::
        \int_{-\infty}^{\infty} e^{-x^2} f(x)\,dx \approx
            \sum_{i=1}^n w_i f(x_i)

    The nodes `x_i` of an order `n` quadrature rule are the roots of `H_n`
    and the weights `w_i` are given by:

    .. math::
        w_i = \frac{2^{n-1} n! \sqrt{\pi}}{n^2 \left(H_{n-1}(x_i)\right)^2}

    Parameters
    ==========

    n : the order of quadrature

    n_digits : number of significant digits of the points and weights to return

    Returns
    =======

    (x, w) : the ``x`` and ``w`` are lists of points and weights as Floats.
             The points `x_i` and weights `w_i` are returned as ``(x, w)``
             tuple of lists.

    Examples
    ========

    >>> from sympy.integrals.quadrature import gauss_hermite
    >>> x, w = gauss_hermite(3, 5)
    >>> x
    [-1.2247, 0, 1.2247]
    >>> w
    [0.29541, 1.1816, 0.29541]

    >>> x, w = gauss_hermite(6, 5)
    >>> x
    [-2.3506, -1.3358, -0.43608, 0.43608, 1.3358, 2.3506]
    >>> w
    [0.00453, 0.15707, 0.72463, 0.72463, 0.15707, 0.00453]

    See Also
    ========

    gauss_legendre, gauss_laguerre, gauss_gen_laguerre, gauss_chebyshev_t, gauss_chebyshev_u, gauss_jacobi, gauss_lobatto

    References
    ==========

    .. [1] https://en.wikipedia.org/wiki/Gauss-Hermite_Quadrature
    .. [2] http://people.sc.fsu.edu/~jburkardt/cpp_src/hermite_rule/hermite_rule.html
    .. [3] http://people.sc.fsu.edu/~jburkardt/cpp_src/gen_hermite_rule/gen_hermite_rule.html
    """
    x = Dummy("x")
    p = hermite_poly(n, x, polys=True)
    p1 = hermite_poly(n-1, x, polys=True)
    xi = []
    w = []
    for r in p.real_roots():
        if isinstance(r, RootOf):
            r = r.eval_rational(S(1)/10**(n_digits+2))
        xi.append(r.n(n_digits))
        w.append(((2**(n-1) * factorial(n) * sqrt(pi)) /
                 (n**2 * p1.subs(x, r)**2)).n(n_digits))
    return xi, w

digits = 36
npoints = list(range(1,17))

for n in npoints:
    xi, wi = gauss_hermite(n, digits)

    print('\nN = %3d'%(n))
    for i in range(0,n):
        print('%3d %40.36f %40.36f'%(i+1,xi[i],wi[i]))

max_stars_repo_name: BryanFlynt/PolyCalc

max_issues_repo_name: BryanFlynt/PolyCalc

max_forks_repo_name: BryanFlynt/PolyCalc

reason: import numpy,import sympy,from sympy


text:
#!/usr/bin/env python
# -*- coding: utf-8 -*-

import click
import numpy as np
import random
import matplotlib.pyplot as plt
from six.moves import cPickle as pickle
import os
import math

'''
tracer-sim --output=$TRACERDIR/data/test --tag=train --num_reads=1000
tracer-sim --output=$TRACERDIR/data/test --tag=test --num_reads=100

tracer-sim --output=$TRACERDIR/data --tag=train --num_reads=1000 --read_length_mean=100 --read_length_max_var=10

'''

@click.command()
@click.option('--output', required=True)
@click.option('--tag', default="train")
@click.option('--read_length_mean', default=20000)
@click.option('--read_length_max_var', default=5000)
@click.option('--num_reads', default=10)
@click.option('--steps_per_base', default=10)
@click.option('--noise_max', default=25)
@click.option('--poly_pause_prob', default=0.9)
@click.option('--poly_stall_prob', default=0.2)
@click.option('--poly_max_pause', default=60)
@click.option('--poly_max_stall', default=30)
@click.option('--mean_intensity_adenine', default=200)
#@click.option('--mean_intensity_adenine_2', default=200)
#@click.option('--mean_intensity_adenine_3', default=200)
@click.option('--mean_intensity_thiamine', default=150)
#@click.option('--mean_intensity_thiamine_2', default=150)
#@click.option('--mean_intensity_thiamine_3', default=150)
@click.option('--mean_intensity_guanine', default=110)
#@click.option('--mean_intensity_guanine_2', default=110)
#@click.option('--mean_intensity_guanine_3', default=110)
@click.option('--mean_intensity_cytosine', default=90)
#@click.option('--mean_intensity_cytosine_2', default=90)
#@click.option('--mean_intensity_cytosine_3', default=90)
@click.option('--base_noise_fraction', default=0.6) # The percentage variation to apply to the intensity of a signal from a base incorporation for noise.
@click.option('--make_figs', default=False)
@click.option('--vocab_size', default=80000)
def cli(output, tag, read_length_mean, read_length_max_var, num_reads,
        steps_per_base, noise_max, poly_pause_prob, poly_stall_prob,
        poly_max_pause, poly_max_stall, mean_intensity_adenine,
        mean_intensity_thiamine, mean_intensity_guanine, mean_intensity_cytosine,
        base_noise_fraction, make_figs, vocab_size):
    """Simulate instrument traces, for development purposes."""
    random.seed(1234)
    noise_max = 0.1

    # So "stalling" vs "pausing" is probably not the right language and this simulator is
    # not going to be at all kinetically accurate. But it will provide well understood input
    # that can be used for development.

    poly_pause_prob = 0.1 # A polymerase pausing after a base incorporation
    poly_stall_prob = 0.1 # A polymerase stalling during base incorporation
    poly_max_pause = 10 # Max duration polymerase can pause
    poly_max_stall = 10 # Max duration polymerase can stall
    mean_intensities = [mean_intensity_adenine, mean_intensity_thiamine, mean_intensity_guanine, mean_intensity_cytosine]
    bases = ["A", "T", "G", "C"]

    # Need to match elsewhere
    #PAD_ID = 0
    #GO_ID = 1
    #EOS_ID = 2
    #UNK_ID = 3
    go_id = 1
    #eos_id = 2
    # Need to shift everything by 4

    if not os.path.exists(output):
        os.mkdir(output)

    raw_trace_fname = os.path.join(output, tag + '.raw.csv')
    raw_seq_fname = os.path.join(output, tag + '.seq.txt')
    encoded_trace_fname = os.path.join(output, tag + '.ids' + str(vocab_size) + '.in')
    encoded_read_fname = os.path.join(output, tag + '.ids' + str(8) + '.out')

    with open(encoded_read_fname, "w") as f_read_encoded:
        with open(encoded_trace_fname, "w") as f_trace_encoded:
            with open(raw_seq_fname, "w") as f_read:
                for i in range(0, num_reads):
                    f_read_encoded.write(str(go_id) + " ")
                    f_trace_encoded.write(str(go_id) + " ")

    #for i in range(0, num_reads):
        #raw_trace_fname = os.path.join(output, tag + '-' + str(i) + '.raw.csv')
        #raw_seq_fname = os.path.join(output, tag + '-' + str(i) + '.seq.txt')
        #encoded_trace_fname = os.path.join(output, tag + '-' + str(i) + '.ids' + str(vocab_size) + '.out')
        #encoded_read_fname = os.path.join(output, tag + '-' + str(i) + '.ids' + str(8) + '.in')
        #with open(encoded_read_fname, "w") as f_read_encoded:
            #f_read_encoded.write("<GO> ")
            #with open(encoded_trace_fname, "w") as f_trace_encoded:
                #f_trace_encoded.write("<GO> ")
                #with open(raw_seq_fname, "w") as f_read:

                    with open(raw_trace_fname, "w") as f:
                        read_length = read_length_mean + random.randint(-1*read_length_max_var, read_length_max_var)
                        for j in range(0, read_length):
                            key = random.randint(0,3)
                            f_read_encoded.write(str(key + 4) + " ")
                            f_read.write(bases[key])
                            mod_val = mean_intensities[key]
                            stall_duration = 0
                            if random.random() < poly_stall_prob:
                                stall_duration = random.randint(0, poly_max_stall)
                            steps = random.randint(1, steps_per_base)
                            for k in range(0, steps + stall_duration):
                                out = gen_noisy_out(noise_max)
                                out[key] += mod_val*(1-base_noise_fraction) + random.random()*mod_val*base_noise_fraction
                                out[key + 4] = 1
                                dump(f, out)
                                encoded = encode_trace_step(out[:4], vocab_size)
                                f_trace_encoded.write(str(encoded) + " ")
                            if random.random() < poly_pause_prob: # Pause the polymerase after incorporation
                                pause_duration = random.randint(0, poly_max_pause)
                                mod_val = 0
                                for i in range(0, pause_duration):
                                    out = gen_noisy_out(noise_max)
                                    out[key + 4] = 1
                                    dump(f, out)
                                    encoded = encode_trace_step(out[:4], vocab_size)
                                    f_trace_encoded.write(str(encoded) + " ")
                        f_read.write("\n")
                    #f_trace_encoded.write(str(eos_id) + "\n")
                    #f_read_encoded.write(str(eos_id) + "\n")
                    f_trace_encoded.write("\n")
                    f_read_encoded.write("\n")

        if make_figs:
            plot_read(fname, fname+'.png')

def gen_noisy_out(noise_max):
    out = [random.uniform(0, noise_max),random.uniform(0, noise_max),random.uniform(0, noise_max),random.uniform(0, noise_max),0,0,0,0]
    return out

def dump(fh, out):
    fh.write(str(out[0]) + ',' + str(out[1]) + ',' + str(out[2]) + ',' +str(out[3]) + ',' +str(out[4]) + ',' +str(out[5]) + ',' + str(out[6]) + ',' + str(out[7]) + '\n')

def plot_read(inpath, outpath):
    max_to_plot = 400
    linecount = 0
    a = np.zeros(max_to_plot)
    t = np.zeros(max_to_plot)
    g = np.zeros(max_to_plot)
    c = np.zeros(max_to_plot)

    with open(inpath, 'r') as f:
        for line in f:
            arr = line.split(',')
            a[linecount] = arr[0]
            t[linecount] = arr[1]
            g[linecount] = arr[2]
            c[linecount] = arr[3]
            linecount += 1
            if linecount >= max_to_plot:
                break

    plt.figure(figsize=(20, 2), dpi=100)
    plt.ylim([0, 400])
    plt.plot(a, 'r', t, 'b', g, 'g', c, 'y')
    plt.savefig(outpath)

max_stars_repo_name: cwbeitel/tracer


text:
import sys
import configparser
import logging.config
from copy import deepcopy

import numpy as np
import pandas as pd
import multiprocess as mp
from sklearn.model_selection import KFold
from sklearn.metrics import average_precision_score, roc_auc_score, r2_score
import xgboost as xgb
import shap

from modeling_utils import *

## Intialize logger
logging.config.fileConfig('logging.ini', disable_existing_loggers=False)
logger = logging.getLogger(__name__)

## Load default configuration
config = configparser.ConfigParser()
config.read('config.ini')
RAND_NUM = int(config['DEFAULT']['rand_num'])
np.random.seed(RAND_NUM)
MAX_RECURSION = int(config['DEFAULT']['max_recursion'])
sys.setrecursionlimit(MAX_RECURSION)
MAX_CV_FOLDS = int(config['DEFAULT']['max_cv_folds'])
BG_GENE_NUM = 1000

class TFPRExplainer:
    def __init__(self, tf_feat_mtx_dict, nontf_feat_mtx, features, label_df_dict):
        self.tfs = np.sort(list(label_df_dict.keys()))
        self.genes = label_df_dict[self.tfs[0]].index.values
        self.feats = features
        self.n_tfs = len(self.tfs)
        self.n_genes = len(self.genes)
        self.k_folds = min(MAX_CV_FOLDS, len(self.tfs))

        self.tg_pairs = [tf + ':' + gene for tf in self.tfs for gene in self.genes]
        self.tf_X = np.vstack([tf_feat_mtx_dict[tf] for tf in self.tfs])
        self.nontf_X = nontf_feat_mtx
        self.y = np.hstack([label_df_dict[tf].values for tf in self.tfs])

    def cross_validate(self):
        """Cross valdiate a classifier or regressor using multiprocessing.
        """
        with mp.Pool(processes=self.k_folds) as pool:
            mp_results = {}
            tf_pseudo_X = np.empty((len(self.tfs), 0))

            kfolds = KFold(n_splits=self.k_folds, shuffle=True, random_state=RAND_NUM)

            for k, (tf_tr_idx, tf_te_idx) in enumerate(kfolds.split(tf_pseudo_X)):
                tr_idx = expand_tf2gene_index(tf_tr_idx, self.n_genes)
                te_idx = expand_tf2gene_index(tf_te_idx, self.n_genes)

                y_tr, y_te = self.y[tr_idx], self.y[te_idx]
                tf_X_tr, tf_X_te = self.tf_X[tr_idx], self.tf_X[te_idx]

                tf_X_tr, tf_X_te = standardize_feat_mtx(tf_X_tr, tf_X_te, 'zscore')
                nontf_X, _ = standardize_feat_mtx(self.nontf_X, None, 'zscore')

                mp_results[k] = pool.apply_async(
                    train_and_predict,
                    args=(
                        k,
                        (tf_X_tr, y_tr),
                        (tf_X_te, y_te),
                        nontf_X,
                        (self.tfs[tf_tr_idx], self.tfs[tf_te_idx]),
                        self.genes))

            self.cv_results = compile_mp_results(mp_results)

    def explain(self):
        """Use SHAP values to features' contributions to predict the
        responsiveness of a gene.
        """
        with mp.Pool(processes=self.k_folds) as pool:
            mp_results = {}

            for k, y_te in enumerate(self.cv_results['preds']):
                n_tfs_te = len(y_te['tf'].unique())
                n_tfs_tr = self.n_tfs - n_tfs_te

                y_te['tf:gene'] = y_te['tf'] + ':' + y_te['gene']
                te_tg_pairs = y_te['tf:gene'].values
                te_idx = [self.tg_pairs.index(tg_pair) for tg_pair in te_tg_pairs]

                tr_idx = sorted(set(range(len(self.tg_pairs))) - set(te_idx))
                logger.info('Explaining {} genes in fold {}'.format(len(te_idx), k))

                tf_X_tr, tf_X_te = self.tf_X[tr_idx], self.tf_X[te_idx]

                tf_X_tr, tf_X_te = standardize_feat_mtx(tf_X_tr, tf_X_te, 'zscore')
                nontf_X, _ = standardize_feat_mtx(self.nontf_X, None, 'zscore')

                X_tr = np.hstack([tf_X_tr, np.vstack([nontf_X for i in range(n_tfs_tr)])])
                X_te = np.hstack([tf_X_te, np.vstack([nontf_X for i in range(n_tfs_te)])])

                bg_idx = np.random.choice(
                    range(X_tr.shape[0]), BG_GENE_NUM, replace=False)
                mp_results[k] = pool.apply_async(
                    calculate_tree_shap,
                    args=(
                        self.cv_results['models'][k],
                        X_te, te_tg_pairs, X_tr[bg_idx],))

            self.shap_vals = [mp_results[k].get() for k in sorted(mp_results.keys())]

    def save(self, dirpath):
        """Save output data.
        """
        pd.concat(self.cv_results['preds']).to_csv(
            '{}/preds.csv.gz'.format(dirpath),
            index=False, compression='gzip')

        pd.concat(self.cv_results['stats']).to_csv(
            '{}/stats.csv.gz'.format(dirpath),
            index=False, compression='gzip')

        np.savetxt(
            '{}/feats.csv.gz'.format(dirpath), np.array(self.feats),
            fmt='%s', delimiter=',')

        np.savetxt(
            '{}/tf_gene_pairs.csv.gz'.format(dirpath), np.array(self.tg_pairs),
            fmt='%s', delimiter=',')

        np.savetxt(
            '{}/feat_mtx_tf.csv.gz'.format(dirpath), self.tf_X,
            fmt='%.8f', delimiter=',')
        np.savetxt(
            '{}/feat_mtx_nontf.csv.gz'.format(dirpath), self.nontf_X,
            fmt='%.8f', delimiter=',')

        # TODO
        for k, df in enumerate(self.shap_vals):
            df['cv'] = k
            self.shap_vals[k] = df
        pd.concat(self.shap_vals).to_csv(
            '{}/feat_shap_wbg.csv.gz'.format(dirpath),
            index=False, compression='gzip')

def train_and_predict(k, D_tr, D_te, nontf_X, tfs, genes):
    """Train classifier and predict gene responses.
    """
    logger.info('Cross validating fold {}'.format(k))

    tf_X_tr, y_tr = D_tr
    tf_X_te, y_te = D_te
    tfs_tr, tfs_te = tfs

    X_tr = np.hstack([tf_X_tr, np.vstack([nontf_X for i in range(len(tfs_tr))])])
    X_te = np.hstack([tf_X_te, np.vstack([nontf_X for i in range(len(tfs_te))])])

    ## Train classifier and test
    model = train_classifier(X_tr, y_tr)

    y_pred = pd.DataFrame(
        data=model.predict_proba(X_te),
        columns=model.classes_)[1].values

    ## Calculate AUC for each TF
    stats_df = pd.DataFrame()
    preds_df = pd.DataFrame()
    n_genes = len(genes)

    for i, tf in enumerate(tfs_te):
        idx = list(range(i * n_genes, (i + 1) * n_genes))
        auprc = average_precision_score(y_te[idx], y_pred[idx])
        auroc = roc_auc_score(y_te[idx], y_pred[idx])

        preds_df = preds_df.append(pd.DataFrame(
            {'gene': genes, 'tf': [tf] * n_genes, 'label': y_te[idx], 'pred': y_pred[idx]}),
            ignore_index=True)
        stats_df = stats_df.append(pd.DataFrame(
            {'cv': [k], 'tf': [tf], 'auroc': [auroc], 'auprc': [auprc]}),
            ignore_index=True)

        logger.info('CV performance for TF {} in fold {}: AUPRC={:.3f}'.format(tf, k, auprc))

    return {'preds': preds_df, 'stats': stats_df, 'models': model}

def train_classifier(X, y):
    """Train a XGBoost classifier.
    """
    model = xgb.XGBClassifier(
        n_estimators=2500,
        learning_rate=.01,
        booster='gbtree',
        gamma=5,
        colsample_bytree=.8,
        subsample=.8,
        n_jobs=-1,
        random_state=RAND_NUM
    )
    model.fit(X, y)
    return model

def train_regressor(X, y):
    """Train a XGBoost regressor.
    """
    model = xgb.XGBRegressor(
        n_estimators=2500,
        learning_rate=.01,
        objective='reg:squarederror',
        booster='gbtree',
        n_jobs=-1,
        random_state=RAND_NUM
    )
    model.fit(X, y)
    return model

def calculate_tree_shap(model, X, genes, X_bg):
    """Calcualte SHAP values for tree-based model.
    """
    n_genes, n_feats = X.shape

    ## Calculate SHAP values
    explainer = shap.TreeExplainer(model, X_bg)
    shap_mtx = explainer.shap_values(X, approximate=False, check_additivity=False)

    ## Convert wide to long format
    shap_df = pd.DataFrame(
        data=shap_mtx,
        index=genes,
        columns=['feat' + str(i) for i in range(n_feats)])
    shap_df.index.name = 'gene'
    shap_df = shap_df.reset_index()
    shap_df = pd.wide_to_long(shap_df, 'feat', 'gene', 'feat_idx').reset_index()
    return shap_df

def expand_tf2gene_index(t, n):
    g = []
    for i in t:
        g += list(range(i * n, (i + 1) * n))
    return g

max_stars_repo_name: wsjung/TFPertRespExplainer

max_issues_repo_name: wsjung/TFPertRespExplainer

max_forks_repo_name: wsjung/TFPertRespExplainer


text:
'''
Arm
move
@version 0.1
@author Juan Cely <juanscelyg@gmail.com>
Research Group of Robots and Intelligent Machines
Date: August/2017
'''
import GRMI_MUR.Arm.arm_parts as Arm_parts
import GRMI_MUR.Arm.kine as Arm_kine
import GRMI_MUR.Simulate.arm as vrep_arm
import math
import numpy as np

def planner_616(x,y,z,theta,phi,Ttol,Tac,n):
	Ttol=float(Ttol)
	Tac=float(Tac)
	if n<10:
		n=10;
	qo=vrep_arm.GetPosition()
	qf=Arm_kine.inversekine(x,y,z,theta,phi)
	dt=Ttol/(n)
	dq=np.empty(len(qo))
	for i in range(0,len(qo)):
		dq[i]=(qf[i]-qo[i])/(n)
	tiempo=np.arange(0.0,Ttol,dt)
	q=np.empty([len(tiempo),len(qo)])
	q[0,:]=qo
	for i in range(1,len(tiempo)):
		q[i,:]=q[i-1,:]+dq
	qp=np.diff(q,axis=0)
	Vmax=np.empty(len(qo))
	for i in range(len(qo)):
		Vmax[i]=dq[i]/dt
	#Soft
	#q soft
	qs=np.empty([len(tiempo),len(qo)])
	qs[0,:]=qo;
	for j in range(0,len(qo)):
		q_o=qo[j]
		for i in range(1,len(tiempo)):
			if tiempo[i]<=Tac:
				t=tiempo[i]
				a0=q_o
				a1=0
				a2=0
				a3=0
				a4=(5/(2*(math.pow(Tac,3))))*Vmax[j]
				a5=(-3/(math.pow(Tac,4)))*Vmax[j]
				a6=(1/(math.pow(Tac,5)))*Vmax[j]
				qs[i,j]=(a6*(math.pow(t,6)))+(a5*(math.pow(t,5)))+(a4*(math.pow(t,4)))+(a3*(math.pow(t,3)))+(a2*(math.pow(t,2)))+(a1*t)+a0
			elif tiempo[i]>=Tac and tiempo[i]<(Ttol-Tac):
				t=tiempo[i]-Tac
				b0=q_o+((Tac/2)*Vmax[j])
				b1=Vmax[j]
				qs[i,j]=(b1*t)+b0
			elif tiempo[i]>=(Ttol-Tac):
				t=tiempo[i]-(Ttol-Tac)
				c0=(Ttol-(3.0/2)*Tac)*Vmax[j]+q_o
				c1=Vmax[j]
				c2=0
				c3=0
				c4=(-5/(2*(math.pow(Tac,3))))*Vmax[j]
				c5=(3/(math.pow(Tac,4)))*Vmax[j]
				c6=(-1/(math.pow(Tac,5)))*Vmax[j]
				qs[i,j]=(c6*(math.pow(t,6)))+(c5*(math.pow(t,5)))+(c4*(math.pow(t,4)))+(c3*(math.pow(t,3)))+(c2*(math.pow(t,2)))+(c1*t)+c0
	# To velocity
	qp_0=np.diff(qs, axis=0)
	limite=np.round(len(qp_0)/2,0)
	qp_1=qp_0[0:limite-1,:]
	qp_2=qp_0[limite:len(qp_0),:]
	qp_s=np.concatenate((qp_1,qp_2), axis=0)
	qp_s=np.insert(qp_s, len(qp_s), [0, 0, 0, 0, 0],axis=0)
	qp_s=np.insert(qp_s, 0, [0, 0, 0, 0, 0],axis=0)
	# To acc
	qpp_0=np.diff(qp_s, axis=0)
	limite=np.round(len(qpp_0)/2,0)
	qpp_1=qpp_0[0:limite-1,:]
	qpp_2=qpp_0[limite:len(qpp_0),:]
	qpp_s=np.concatenate((qpp_1,qpp_2), axis=0)
	qpp_s=np.insert(qpp_s, len(qpp_s), [0, 0, 0, 0, 0],axis=0)
	qpp_s=np.insert(qpp_s, 0, [0, 0, 0, 0, 0],axis=0)
	return np.round(qs,4), np.round(qp_s,4), np.round(qpp_s,4), np.round(tiempo,4)


text:
# Trajectory optimised Linear Quadratic Regulator with replans Sims
# Motion planning for a robot in a fully observable environment under uncertain process noise.
# Author: Mohamed Naveed Gul Mohamed
# email:mohdnaveed96@gmail.com
# Date: Jan 23rd 2020
import casadi as c
import math as m
import matplotlib
matplotlib.use('Agg')

from matplotlib import cm
from matplotlib.patches import Rectangle
from matplotlib.patches import Ellipse
import numpy as np
from numpy import linalg
import pylab
import sys,os
import time

import Algorithm_classes as algo
from robot_models import youBot_model
from robot_models import car_model
from robot_models import car_w_trailers
import simulation_params as params

no_iters = 100

PLOT = False #True #
PLOT_ROBOT_BOUNDARY = True #False #
FILE_WRITE = False #True #
HPRC = False
Model = 'car_w_trailers' #'car_model' #'youBot_model'
OL_time = 0 #Open loop time taken.

Xg = np.array([2,2,0,0,0,0])

if FILE_WRITE:
    filename = "/home/naveed/Dropbox/Research/Data/WAFR20/car_w_trailers/TLQR2_1_wo_obs_modified.csv"
    file = open(filename,"a")
    file.write('epsilon' + ',' + 'Average Cost' + ',' + 'Cost variance' + ',' + 'Average Time' + ',' + 'Average Replans' + ','
        + 'Replan Variance' + ',' + 'replan_bound' + '\n' )

if Model == 'car_w_trailers':
    robot = car_w_trailers(dt=0.1)
    #open-loop optimisation gain
    R = c.DM([[5,0],[0,5]])
    Q = c.DM([[10,0,0,0,0,0],[0,10,0,0,0,0],[0,0,1,0,0,0],[0,0,0,1,0,0],[0,0,0,0,1,0],[0,0,0,0,0,1]])
    Qf = c.DM([[50,0,0,0,0,0],[0,50,0,0,0,0],[0,0,50,0,0,0],[0,0,0,5,0,0],[0,0,0,0,5,0],[0,0,0,0,0,5]])

def blockPrint():
    sys.stdout = open(os.devnull, 'w')

# Restore
def enablePrint():
    sys.stdout = sys.__stdout__

def calc_cost(U, X, T):

    Cost_OL_vec = np.zeros(T+1)
    Cost_OL = 0

    for i in range(T):

        Cost_OL = Cost_OL + c.mtimes(c.mtimes(c.reshape(X[:,i],(1,robot.nx)),Q),c.reshape(X[:,i],(robot.nx,1))) + \
                            c.mtimes(c.mtimes(c.reshape(U[:,i],(1,robot.nu)),R),c.reshape(U[:,i],(robot.nu,1)))

        if params.OBSTACLES:
            Obstacle_cost = control.obstacle_cost_func(X_path_ss[0:2,i])
            Cost_OL = Cost_OL + Obstacle_cost

        Cost_OL_vec[i] = Cost_OL

    Cost_OL = Cost_OL + c.mtimes(c.mtimes(c.reshape(Xg - X[:,T],(1,robot.nx)),Qf),c.reshape(Xg - X[:,T],(robot.nx,1)))
    Cost_OL_vec[T] = Cost_OL

    return Cost_OL_vec

def SSP(T, X0, Q, R, Qf, U_guess):
    #State space planning

    control = algo.SSP(T, X0, Xg, R, Q, Qf, params.Xmin, params.Xmax)

    #U_guess = np.zeros((robot.nu,T))

    U_guess = np.reshape(U_guess.T,(robot.nu*T,1))

    print("Solving state space problem...")
    U_opti_ss = control.solve_SSP(robot,X0, np.zeros(robot.nu), U_guess)
    print("Done.")
    #print("Uss:",U_opti_ss)
    X_path_ss = np.zeros((robot.nx,T+1))
    X_path_ss[:,0] = X0
    Cost_OL_vec = np.zeros(T+1)
    Cost_OL = 0

    for i in range(T):

        X_path_ss[:,i+1] = np.reshape(robot.kinematics(X_path_ss[:,i],U_opti_ss[:,i]),(robot.nx,)) #to make it compatible in dimension reshaping from (3,1) to 3

        Cost_OL = Cost_OL + c.mtimes(c.mtimes(c.reshape(X_path_ss[:,i],(1,robot.nx)),Q),c.reshape(X_path_ss[:,i],(robot.nx,1))) + \
                            c.mtimes(c.mtimes(c.reshape(U_opti_ss[:,i],(1,robot.nu)),R),c.reshape(U_opti_ss[:,i],(robot.nu,1)))

        if params.OBSTACLES:
            Obstacle_cost = control.obstacle_cost_func(X_path_ss[0:2,i])
            Cost_OL = Cost_OL + Obstacle_cost

        Cost_OL_vec[i] = Cost_OL

    Cost_OL = Cost_OL + c.mtimes(c.mtimes(c.reshape(Xg - X_path_ss[:,T],(1,robot.nx)),Qf),c.reshape(Xg - X_path_ss[:,T],(robot.nx,1)))
    Cost_OL_vec[T] = Cost_OL

    print "U nom:", U_opti_ss

    return U_opti_ss, X_path_ss, Cost_OL_vec

def TLQR2_run_varying_epsilon():

    global OL_time
    T = 35
    X0 = np.array([0,0,m.pi/3,0,0,0])

    #LQR gain
    Wx = c.DM([[10,0,0,0,0,0],[0,10,0,0,0,0],[0,0,1,0,0,0],[0,0,0,1,0,0],[0,0,0,0,1,0],[0,0,0,0,0,1]]) #10
    Wxf = c.DM([[50,0,0,0,0,0],[0,50,0,0,0,0],[0,0,1,0,0,0],[0,0,0,1,0,0],[0,0,0,0,1,0],[0,0,0,0,0,1]])  #100
    Wu = c.DM([[1,0],[0,1]])#10

    U_guess = np.zeros((robot.nu,T))

    OL_start = time.time()
    U_opti, X_nom, Cost_nom = SSP(T,X0,Q,R,Qf,U_guess)

    control = algo.LQR(Wx,Wu,Wxf,T)

    print("Solving for LQR gain...")
    K_lqr = control.solve_K(robot, X_nom, U_opti)

    OL_end = time.time()

    OL_time = OL_end - OL_start

    replan_bound_vec = np.array([.02])
    print "replan bound range:", replan_bound_vec

    epsi_range = np.linspace(0.21,.40,20)
    print "Epsilon range:", epsi_range

    blockPrint()
    for replan_bound in replan_bound_vec:
        for epsilon in epsi_range:
            TLQR2(control, T, X0, U_opti, X_nom, K_lqr, Cost_nom, replan_bound, epsilon)

def TLQR2(control, T, X0, U_opti_ini, X_nom_ini, K_lqr_ini, Cost_nom_ini, replan_bound, epsilon):

    print("Executing Plan...")
    #Execution of plan

    start_exec = time.time()

    Cost_CL_iter = np.zeros(no_iters)
    Replan_iter = np.zeros(no_iters) #vector to store # of replans in every iteration.

    iter = 0
    #np.random.seed(4)
    while iter < no_iters:

        U_opti = U_opti_ini
        X_nom = X_nom_ini
        K_lqr = K_lqr_ini
        Cost_nom = Cost_nom_ini

        #try:

        X_act_CL = np.zeros((robot.nx,T+1))

        X_act_CL[:,0] = X0 #initial condition

        U_CL = np.zeros((robot.nu,T))
        Cost_CL = 0
        replan_count = 0
        replan = False
        replan_index = []

        #closed loop implementation
        for i in range(T):
            #Applying LQR control

            #print "X nom:", X_nom[:,i], " X act:", X_act_CL[:,i]
            if replan:
                replan_count += 1

                replan_index.append(i)

                U_guess = U_opti[:,i:]
                print "Replanning at step:", i
                U_opti[:,i:T], X_nom[:,i:T+1],_ = SSP(T-i,X_act_CL[:,i],Q,R,Qf, U_guess) #calculating new nominal

                #print "Cost_nom:", calc_cost(U_opti,X_nom,T)

                control.T = T - i #setting new horizon.
                K_lqr[:,i:T+1] = control.solve_K(robot, X_nom[:,i:T+1], U_opti[:,i:T]) #calculating new LQR gains

                replan = False

            U_temp = U_opti[:,i] - c.mtimes(c.reshape(K_lqr[:,i],robot.nu,robot.nx),(X_act_CL[:,i] - X_nom[:,i]))

            U_CL[:,i] = np.reshape(control.U_bounds(robot,U_temp),(robot.nu,))

            X_act_CL[:,i+1] = np.reshape(robot.kinematics(X_act_CL[:,i],U_CL[:,i],epsilon),(robot.nx,))

            Cost_CL = Cost_CL + c.mtimes(c.mtimes(c.reshape(X_act_CL[:,i],(1,robot.nx)),Q),c.reshape(X_act_CL[:,i],(robot.nx,1))) + \
                                c.mtimes(c.mtimes(c.reshape(U_CL[:,i],(1,robot.nu)),R),c.reshape(U_CL[:,i],(robot.nu,1)))

            if params.OBSTACLES:
                Obstacle_cost = control.obstacle_cost_func(X_act_CL[0:2,i])
                Cost_CL = Cost_CL + Obstacle_cost

            #enablePrint()
            print "Cost_CL:", Cost_CL, " Cost_nom:", Cost_nom[i], " deviation:", abs(Cost_CL - Cost_nom[i])/Cost_nom[i]
            blockPrint()

            if abs(Cost_CL - Cost_nom[i])/Cost_nom[i] > replan_bound:
                replan = True

        Cost_CL = Cost_CL + c.mtimes(c.mtimes(c.reshape(Xg - X_act_CL[:,T],(1,robot.nx)),Qf),c.reshape(Xg - X_act_CL[:,T],(robot.nx,1)))
        #enablePrint()
        print "Cost_CL:", Cost_CL, " Cost_nom:", Cost_nom[T], " deviation:", abs(Cost_CL - Cost_nom[T])/Cost_nom[T]
        blockPrint()

        Cost_CL_iter[iter] = Cost_CL
        Replan_iter[iter] = replan_count
        iter = iter + 1

        # except:
        #         print "Unexpected error:", sys.exc_info()

    end_exec = time.time()
    time_taken = OL_time + (end_exec - start_exec)/no_iters

    enablePrint()
    print "epsilon:",epsilon , "   Average Cost:",np.mean(Cost_CL_iter), "  Cost var:", np.var(Cost_CL_iter), " Average replans:", np.mean(Replan_iter), " Replan Var:", np.var(Replan_iter), "replan_bound",replan_bound," Time taken:",time_taken
    blockPrint()

    if FILE_WRITE:
        file.write(str(epsilon) + ',' + str(np.mean(Cost_CL_iter)) + ',' +
            str(np.var(Cost_CL_iter)) + ',' + str(time_taken) + ',' + str(np.mean(Replan_iter)) + ','
            + str(np.var(Replan_iter))+ ',' + str(replan_bound) + '\n')

    if PLOT:
        pylab.figure(1)
        pylab.plot(X_nom[0,:],X_nom[1,:],'y',linewidth=2,label="Nominal")
        pylab.plot(X_act_CL[0,:],X_act_CL[1,:],'--b',linewidth=2,label="T-LQR2")
        pylab.plot(X0[0],X0[1], 'co',label='Start',markersize=10)
        pylab.plot(Xg[0], Xg[1], 'g*',label='Goal',markersize=10)

        if PLOT_ROBOT_BOUNDARY:
            #plotting robot boundary
            diag = m.sqrt(robot.length**2 + robot.breadth**2)/2
            alp = m.atan2(robot.breadth,robot.length)
            ax = pylab.gca()

            for i in np.arange(0,T,10):
                ax.add_patch(Rectangle((X_act_CL[0,i] - diag*m.cos(X_act_CL[2,i] + alp),
                    X_act_CL[1,i] - diag*m.sin(X_act_CL[2,i] + alp)), robot.length, robot.breadth,angle=m.degrees(X_act_CL[2,i]),fill=None,linewidth=.1,color='b'))

            ax.add_patch(Rectangle((X_act_CL[0,T] - diag*m.cos(X_act_CL[2,T] + alp),
                X_act_CL[1,T] - diag*m.sin(X_act_CL[2,T] + alp)), robot.length, robot.breadth,angle=m.degrees(X_act_CL[2,T]),fill=None,linewidth=.1,color='b'))

        pylab.xlabel('X (meters)')
        pylab.ylabel('Y (meters)')
        legend = pylab.legend(loc=0)
        frame = legend.get_frame()

        pylab.show()

def TLQR2_hprc(epsilon,replan_bound):

    T = 40
    X0 = np.array([0,0,m.pi/3,0,0,0])

    #LQR gain
    Wx = c.DM([[10,0,0,0,0,0],[0,10,0,0,0,0],[0,0,1,0,0,0],[0,0,0,1,0,0],[0,0,0,0,1,0],[0,0,0,0,0,1]]) #10
    Wxf = c.DM([[50,0,0,0,0,0],[0,50,0,0,0,0],[0,0,1,0,0,0],[0,0,0,1,0,0],[0,0,0,0,1,0],[0,0,0,0,0,1]])  #100
    Wu = c.DM([[1,0],[0,1]]) #10

    U_guess = np.zeros((robot.nu,T))

    OL_start = time.time()
    U_opti, X_nom, Cost_nom = SSP(T,X0,Q,R,Qf,U_guess)

    control = algo.LQR(Wx,Wu,Wxf,T)

    print("Solving for LQR gain...")
    K_lqr = control.solve_K(robot, X_nom, U_opti)

    OL_end = time.time()

    OL_time = OL_end - OL_start

    print("Executing Plan...")
    #Execution of plan

    start_exec = time.time()
    COMPLETE = False
    while not COMPLETE:

        try:

            X_act_CL = np.zeros((robot.nx,T+1))

            X_act_CL[:,0] = X0 #initial condition

            U_CL = np.zeros((robot.nu,T))
            Cost_CL = 0
            replan_count = 0
            replan = False
            replan_index = []

            #closed loop implementation
            for i in range(T):
                #Applying LQR control

                #print "X nom:", X_nom[:,i], " X act:", X_act_CL[:,i]
                if replan:
                    replan_count += 1

                    replan_index.append(i)

                    U_guess = U_opti[:,i:]
                    print "Replanning at step:", i
                    U_opti[:,i:T], X_nom[:,i:T+1],_ = SSP(T-i,X_act_CL[:,i],Q,R,Qf, U_guess) #calculating new nominal

                    #print "Cost_nom:", calc_cost(U_opti,X_nom,T)

                    control.T = T - i #setting new horizon.
                    K_lqr[:,i:T+1] = control.solve_K(robot, X_nom[:,i:T+1], U_opti[:,i:T]) #calculating new LQR gains

                    replan = False

                U_temp = U_opti[:,i] - c.mtimes(c.reshape(K_lqr[:,i],robot.nu,robot.nx),(X_act_CL[:,i] - X_nom[:,i]))

                U_CL[:,i] = np.reshape(control.U_bounds(robot,U_temp),(robot.nu,))

                X_act_CL[:,i+1] = np.reshape(robot.kinematics(X_act_CL[:,i],U_CL[:,i],epsilon),(robot.nx,))

                Cost_CL = Cost_CL + c.mtimes(c.mtimes(c.reshape(X_act_CL[:,i],(1,robot.nx)),Q),c.reshape(X_act_CL[:,i],(robot.nx,1))) + \
                                    c.mtimes(c.mtimes(c.reshape(U_CL[:,i],(1,robot.nu)),R),c.reshape(U_CL[:,i],(robot.nu,1)))

                if params.OBSTACLES:
                    Obstacle_cost = control.obstacle_cost_func(X_act_CL[0:2,i])
                    Cost_CL = Cost_CL + Obstacle_cost

                #enablePrint()
                print "Cost_CL:", Cost_CL, " Cost_nom:", Cost_nom[i], " deviation:", abs(Cost_CL - Cost_nom[i])/Cost_nom[i]
                blockPrint()

                if abs(Cost_CL - Cost_nom[i])/Cost_nom[i] > replan_bound:
                    replan = True

            Cost_CL = Cost_CL + c.mtimes(c.mtimes(c.reshape(Xg - X_act_CL[:,T],(1,robot.nx)),Qf),c.reshape(Xg - X_act_CL[:,T],(robot.nx,1)))
            #enablePrint()
            print "Cost_CL:", Cost_CL, " Cost_nom:", Cost_nom[T], " deviation:", abs(Cost_CL - Cost_nom[T])/Cost_nom[T]
            blockPrint()

            COMPLETE = True

        except:
            print "Unexpected error:", sys.exc_info()

    end_exec = time.time()
    time_taken = OL_time + (end_exec - start_exec)

    return(Cost_CL,time_taken, replan_count)

if __name__=='__main__':

    # if Model == 'car_w_trailers':
    #     #open-loop optimisation gain
    #     R = c.DM([[20,0],[0,10]])
    #     Q = 5*c.DM.eye(robot.nx)
    #     Qf = 1000*c.DM.eye(robot.nx)
    #
    #
    # U_opti, X_nom = SSP(Q,R,Qf)
    #
    # TLQR(U_opti, X_nom, epsilon=0.3)

    TLQR2_run_varying_epsilon()

    # if FILE_WRITE:
    #     file.close()

reason: import numpy,from numpy


text:
import logging
import importlib
import torch

import mne

importlib.reload(logging) # see https://stackoverflow.com/a/21475297/1469195
log = logging.getLogger()
log.setLevel('INFO')
import sys
logging.basicConfig(format='%(asctime)s %(levelname)s : %(message)s',
                     level=logging.INFO, stream=sys.stdout)
from EEGNAS import global_vars
from EEGNAS_experiment import set_params_by_dataset
from EEGNAS.data_preprocessing import get_train_val_test, get_ch_names
from braindecode.torch_ext.util import set_random_seeds
from EEGNAS_experiment import get_normal_settings, parse_args, get_configurations
from naiveNAS import NaiveNAS
from EEGNAS.models_generation import target_model
from braindecode.torch_ext.util import np_to_var
import numpy as np
from torch import nn
from braindecode.torch_ext.util import var_to_np
import torch as th
import matplotlib
matplotlib.use('qt5agg')

args = parse_args(['-e', 'tests', '-c', '../configurations/config.ini'])
global_vars.init_config(args.config)
configs = get_configurations(args.experiment)
assert (len(configs) == 1)
global_vars.set_config(configs[0])

subject_id = 1
low_cut_hz = 0
fs = 250
valid_set_fraction = 0.2
dataset = 'BCI_IV_2a'
data_folder = '../data/'
global_vars.set('dataset', dataset)
set_params_by_dataset()
global_vars.set('cuda', True)
model_select = 'deep4'
model_dir = '143_x_evolution_layers_cross_subject'
model_name = 'best_model_9_8_6_7_2_1_3_4_5.th'
train_set = {}
val_set = {}
test_set = {}
train_set[subject_id], val_set[subject_id], test_set[subject_id] =\
            get_train_val_test(data_folder, subject_id)

# Set if you want to use GPU
# You can also use torch.cuda.is_available() to determine if cuda is available on your machine.
cuda = True
set_random_seeds(seed=20170629, cuda=cuda)

# This will determine how many crops are processed in parallel
input_time_length = 450
# final_conv_length determines the size of the receptive field of the ConvNet
models = {'evolution': torch.load(f'../models/{model_dir}/{model_name}'),
         'deep4': target_model('deep')}
model = models[model_select]
input_time_length = global_vars.get('input_time_len')
stop_criterion, iterator, loss_function, monitors = get_normal_settings()
naiveNAS = NaiveNAS(iterator=iterator, exp_folder=None, exp_name=None,
                    train_set=train_set, val_set=val_set, test_set=test_set,
                    stop_criterion=stop_criterion, monitors=monitors, loss_function=loss_function,
                    config=global_vars.config, subject_id=1, fieldnames=None, strategy='per_subject',
                    csv_file=None, evolution_file=None)
naiveNAS.train_and_evaluate_model(model)

train_batches = list(iterator.get_batches(train_set[subject_id], shuffle=False))
train_X_batches = np.concatenate(list(zip(*train_batches))[0])

new_model = nn.Sequential()
for name, module in model.named_children():
    if 'softmax' in name: break
    new_model.add_module(name, module)

new_model.eval()
pred_fn = lambda x: var_to_np(th.mean(new_model(np_to_var(x).cuda())[:, :, :, 0], dim=2, keepdim=False))
from braindecode.visualization.perturbation import compute_amplitude_prediction_correlations
amp_pred_corrs = compute_amplitude_prediction_correlations(pred_fn, train_X_batches, n_iterations=12,
                                         batch_size=30)

freqs = np.fft.rfftfreq(train_X_batches.shape[2], d=1.0/fs)

alpha_band = {'start': 7, 'stop': 14}
beta_band = {'start': 14, 'stop': 31}
high_gamma_band = {'start': 71, 'stop': 91}
bands = [alpha_band, beta_band, high_gamma_band]

for band in bands:
    band['i_start'] = np.searchsorted(freqs, band['start'])
    band['i_stop'] = np.searchsorted(freqs, band['stop']) + 1
    band['freq_corr'] = np.mean(amp_pred_corrs[:, band['i_start']:band['i_stop']], axis=1)

from braindecode.datasets.sensor_positions import get_channelpos, CHANNEL_10_20_APPROX

ch_names = ['Fz', 'FC3', 'FC1', 'FCz', 'FC2', 'FC4', 'C5', 'C3', 'C1', 'Cz', 'C2', 'C4', 'C6',
            'CP3', 'CP1', 'CPZ', 'CP2', 'CP4', 'P1', 'Pz', 'P3', 'POz']
positions = [get_channelpos(name, CHANNEL_10_20_APPROX) for name in ch_names]
positions = np.array(positions)

import matplotlib.pyplot as plt
from matplotlib import cm
max_abs_val = np.max([np.abs(band['freq_corr']) for band in bands])

fig, axes = plt.subplots(len(bands), global_vars.get('n_classes'))
class_names = ['Left Hand', 'Right Hand', 'Feet', 'Tongue']
for band_i, band in enumerate(bands):
    for i_class in range(global_vars.get('n_classes')):
        ax = axes[band_i, i_class]
        mne.viz.plot_topomap(band['freq_corr'][:, i_class], positions,
                         vmin=-max_abs_val, vmax=max_abs_val, contours=0,
                        cmap=cm.coolwarm, axes=ax, show=False);
        ax.set_title(f'{band["start"]}-{band["stop"]}Hz, {class_names[i_class]}')

plt.show()

max_stars_repo_name: puzis/OverflowPrediction

max_issues_repo_name: puzis/OverflowPrediction


text:
import numpy as np
from scipy import stats

def evaluate_counts_error(
        Operator,
        counts,
        ):
    sample_means = []
    samples =
    pass

    def evalate(a)

def evaluate_error(
        self,
        numberOfSamples=256, # of times to repeat
        sample_size=1024, # number of counts in sample
        ci=0.90, #target CI#,
        f=None,
        replace=False,
        spin_alt=False
        ):
    print('Samples: {}'.format(numberOfSamples))
    print('Sample size: {}'.format(sample_size))
    count_list = []
    N = self.qs.Ns
    if sample_size>=N*8:
        sample_size=int(N/8)
    samplesSD = []
    sample_means = []
    counts_list = {}
    for pauli,counts in self.counts.items():
        count_list = []
        for k,v in counts.items():
            count_list = count_list+[k]*v
        counts_list[pauli]=count_list
    for t in range(numberOfSamples):
        t1 = dt()
        sample_mean  = f(
                self.getRandomRDMFromCounts(
                    counts_list,sample_size
                    )
                )
        if np.isnan(sample_mean):
            continue
        else:
            sample_means.append(sample_mean)
        t2 = dt()
        #print('Time: {}'.format(t2-t1))
    t = stats.t.ppf(ci,N)
    std_err = np.std(np.asarray(sample_means),axis=0) #standard error of mean
    ci = std_err*np.sqrt(sample_size/N)*t
    return ci

def getRandomRDMFromCounts(self,counts_list,length):
    random_counts = {}
    for pauli,clist in counts_list.items():
        random_counts[pauli]={}
        sample_list = np.random.choice(clist,length,replace=False)
        for j in sample_list:
            try:
                random_counts[pauli][j]+=1
            except KeyError:
                random_counts[pauli][j]=1
    #print('Build random list: {}'.format(t3-t5))
    del self.rdm
    self.counts = random_counts
    self.construct()
    #new = self._build_mod_2RDM(random_counts)
    #print('Build 2rdm: {}'.format(t4-t3))
    return self.rdm

max_stars_repo_name: damazz/HQCA

max_issues_repo_name: damazz/HQCA

max_forks_repo_name: damazz/HQCA


text:
# -*- coding: UTF-8 -*-
from sklearn.metrics import classification_report
import csv
import numpy as np
from numpy import matrix
from math import pow
from collections import namedtuple
import math
import random
import os
import json
from sklearn.multiclass import OneVsOneClassifier
from sklearn.svm import LinearSVC
from sklearn.externals import joblib
import pandas as pd
import warnings

from sklearn import svm,preprocessing # new added 1

warnings.filterwarnings("ignore")
class MyModel:
    # 保存神经网络的文件路径
    NN_FILE_PATH = "/home/chang/MagicMirror/server/data/model.m"
    def __init__(self, data_matrix, data_labels, use_file=True):      #  new added  2 初始化的矩阵必须是casia特征集合才能和加载的model匹配
        # 决定了要不要导入nn.json
        self._use_file = use_file
        # 数据集
        self.data_matrix = data_matrix
        self.data_labels = data_labels
        self.scaler = preprocessing.StandardScaler().fit(data_matrix)   # new added  3  对应于casia特征集的缩放器
        self.data_matrix = self.scaler.transform(data_matrix)           # new added 4

        if (not os.path.isfile(self.NN_FILE_PATH) or not use_file):
            # 初始化神经网络
            # 训练并保存
            self.clf	 = svm.SVC(kernel='rbf',C=10,gamma=0.01,cache_size=4000,max_iter=10000,shrinking=True)   # new added 5 将原来的linearSVC 改为 SVC
            ids = np.arange(len(data_labels))

            np.random.shuffle(ids)
            # shuffle

            X = data_matrix[ids]

            y = data_labels[ids]
            self.train(X,y)

            self.save()
        else:
            # 如果nn.json存在则加载
            self._load()

    def train(self, data_matrix,data_labels):
        data_matrix = self.scaler.transform(data_matrix)   # new added 6  缩放特征
        self.clf.fit(data_matrix,data_labels)
        print "training..."
        y_pred = self.predict(data_matrix)
        print(classification_report(data_labels, y_pred))

    def predict(self, test):
        test = self.scaler.transform(test)    # new added  7 缩放特征
        a  =  self.clf.predict(test)
        a = ''.join(a)
        return a
    def save(self):
        if not self._use_file:
            return
        joblib.dump(self.clf, self.NN_FILE_PATH)

    def _load(self):
        if not self._use_file:
            return
        self.clf = joblib.load(self.NN_FILE_PATH)


text:
# -*- coding: utf-8 -*-
"""ctc_pytorch.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1il6A-3w1U_YOPN3xHLWm54yj3iapNqpW

# 导入必要的库

我们需要导入一个叫 [captcha](https://github.com/lepture/captcha/) 的库来生成验证码。

我们生成验证码的字符由数字和大写字母组成。

```sh
pip install captcha numpy matplotlib torch torchvision tqdm
```
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torchvision.transforms.functional import to_tensor, to_pil_image

from captcha.image import ImageCaptcha
from tqdm import tqdm
import random
import numpy as np
from collections import OrderedDict
import string
import argparse
print("Tips: you should run multi-gpu training code using follow command:/n"+"python -m torch.distributed.launch --nproc_per_node=NUM_GPU_YOU_HAS ctc_pytorch_multigpu.py")

parser = argparse.ArgumentParser()
torch.distributed.init_process_group(backend="nccl")
parser.add_argument("--local_rank", type=int)
args = parser.parse_args()
torch.cuda.set_device(args.local_rank)

characters = "-" + string.digits + string.ascii_letters
width, height, n_len, n_classes = 192, 64, 4, len(characters)
n_input_length = 12
print(characters, width, height, n_len, n_classes)

"""# 搭建数据集"""

class CaptchaDataset(Dataset):
    def __init__(self, characters, length, width, height, input_length, label_length):
        super(CaptchaDataset, self).__init__()
        self.characters = characters
        self.length = length
        self.width = width
        self.height = height
        self.input_length = input_length
        self.label_length = label_length
        self.n_class = len(characters)
        self.generator = ImageCaptcha(width=width, height=height)

    def __len__(self):
        return self.length

    def __getitem__(self, index):
        random_str = "".join(
            [random.choice(self.characters[1:]) for j in range(self.label_length)]
        )
        image = to_tensor(self.generator.generate_image(random_str))
        target = torch.tensor(
            [self.characters.find(x) for x in random_str], dtype=torch.long
        )
        input_length = torch.full(
            size=(1,), fill_value=self.input_length, dtype=torch.long
        )
        target_length = torch.full(
            size=(1,), fill_value=self.label_length, dtype=torch.long
        )
        return image, target, input_length, target_length

"""# 初始化数据集生成器"""

batch_size = 128
train_set = CaptchaDataset(
    characters, 1000 * batch_size, width, height, n_input_length, n_len
)
valid_set = CaptchaDataset(
    characters, 100 * batch_size, width, height, n_input_length, n_len
)
train_sampler = torch.utils.data.distributed.DistributedSampler(train_set)
valid_sampler = torch.utils.data.distributed.DistributedSampler(valid_set)
train_loader = DataLoader(
    train_set,
    batch_size=batch_size,
    num_workers=64,
    pin_memory=True,
    sampler=train_sampler,
)
valid_loader = DataLoader(
    valid_set,
    batch_size=batch_size,
    num_workers=12,
    pin_memory=True,
    sampler=valid_sampler,
)

"""# 搭建模型"""

class Model(nn.Module):
    def __init__(self, n_classes, input_shape=(3, 64, 128)):
        super(Model, self).__init__()
        self.input_shape = input_shape
        channels = [32, 64, 128, 256, 256]
        layers = [2, 2, 2, 2, 2]
        kernels = [3, 3, 3, 3, 3]
        pools = [2, 2, 2, 2, (2, 1)]
        modules = OrderedDict()

        def cba(name, in_channels, out_channels, kernel_size):
            modules[f"conv{name}"] = nn.Conv2d(
                in_channels,
                out_channels,
                kernel_size,
                padding=(1, 1) if kernel_size == 3 else 0,
            )
            modules[f"bn{name}"] = nn.BatchNorm2d(out_channels)
            modules[f"relu{name}"] = nn.ReLU(inplace=True)

        last_channel = 3
        for block, (n_channel, n_layer, n_kernel, k_pool) in enumerate(
            zip(channels, layers, kernels, pools)
        ):
            for layer in range(1, n_layer + 1):
                cba(f"{block+1}{layer}", last_channel, n_channel, n_kernel)
                last_channel = n_channel
            modules[f"pool{block + 1}"] = nn.MaxPool2d(k_pool)
        modules[f"dropout"] = nn.Dropout(0.25, inplace=True)

        self.cnn = nn.Sequential(modules)
        self.lstm = nn.LSTM(
            input_size=self.infer_features(),
            hidden_size=128,
            num_layers=2,
            bidirectional=True,
        )
        self.fc = nn.Linear(in_features=256, out_features=n_classes)

    def infer_features(self):
        x = torch.zeros((1,) + self.input_shape)
        x = self.cnn(x)
        x = x.reshape(x.shape[0], -1, x.shape[-1])
        return x.shape[1]

    def forward(self, x):
        x = self.cnn(x)
        x = x.reshape(x.shape[0], -1, x.shape[-1])
        x = x.permute(2, 0, 1)
        x, _ = self.lstm(x)
        x = self.fc(x)
        return x.permute(1, 0, 2)

"""## 测试模型输出尺寸"""

model = Model(n_classes, input_shape=(3, height, width))
model = model.cuda()
model = torch.nn.parallel.DistributedDataParallel(
    model, device_ids=[args.local_rank], output_device=args.local_rank
)

"""# 解码函数和准确率计算函数"""

def decode(sequence):
    a = "".join([characters[x] for x in sequence])
    s = "".join(
        [x for j, x in enumerate(a[:-1]) if x != characters[0] and x != a[j + 1]]
    )
    if len(s) == 0:
        return ""
    if a[-1] != characters[0] and s[-1] != a[-1]:
        s += a[-1]
    return s

def decode_target(sequence):
    return "".join([characters[x] for x in sequence]).replace(" ", "")

def calc_acc(target, output):
    output_argmax = output.detach().permute(1, 0, 2).argmax(dim=-1)
    target = target.cpu().numpy()
    output_argmax = output_argmax.cpu().numpy()
    a = np.array(
        [
            decode_target(true) == decode(pred)
            for true, pred in zip(target, output_argmax)
        ]
    )
    return a.mean()

"""# 训练模型"""

def train(model, optimizer, scheduler, epoch, dataloader):
    model.train()
    loss_mean = 0
    acc_mean = 0
    with tqdm(dataloader) as pbar:
        for batch_index, (data, target, input_lengths, target_lengths) in enumerate(
            pbar
        ):
            data, target = data.cuda(), target.cuda()

            optimizer.zero_grad()
            output = model(data)

            output_log_softmax = F.log_softmax(output, dim=-1)
            output_log_softmax = output_log_softmax.permute(1, 0, 2)
            loss = F.ctc_loss(output_log_softmax, target, input_lengths, target_lengths)

            loss.backward()
            optimizer.step()
            scheduler.step()

            loss = loss.item()
            acc = calc_acc(target, output)

            if batch_index == 0:
                loss_mean = loss
                acc_mean = acc

            loss_mean = 0.1 * loss + 0.9 * loss_mean
            acc_mean = 0.1 * acc + 0.9 * acc_mean

            for param_group in optimizer.param_groups:
                pbar.set_description(
                    f"Epoch: {epoch} Loss: {loss_mean:.4f} Acc: {acc_mean:.4f}   Lr: "
                    + str(param_group["lr"])
                )
                break

def valid(model, optimizer, epoch, dataloader):
    model.eval()
    with tqdm(dataloader) as pbar, torch.no_grad():
        loss_sum = 0
        acc_sum = 0
        for batch_index, (data, target, input_lengths, target_lengths) in enumerate(
            pbar
        ):
            data, target = data.cuda(), target.cuda()

            output = model(data)
            output_log_softmax = F.log_softmax(output, dim=-1)
            output_log_softmax = output_log_softmax.permute(1, 0, 2)
            loss = F.ctc_loss(output_log_softmax, target, input_lengths, target_lengths)

            loss = loss.item()
            acc = calc_acc(target, output)

            loss_sum += loss
            acc_sum += acc

            loss_mean = loss_sum / (batch_index + 1)
            acc_mean = acc_sum / (batch_index + 1)

            pbar.set_description(
                f"Test : {epoch} Loss: {loss_mean:.4f} Acc: {acc_mean:.4f} "
            )

optimizer = torch.optim.Adam(model.parameters(), 1e-3, amsgrad=True)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=100, gamma=0.98)
epochs = 30
for epoch in range(1, epochs + 1):
    train(model, optimizer, scheduler, epoch, train_loader)
    valid(model, optimizer, epoch, valid_loader)
    torch.save(model, "ctc.pth")