text:
"""Proximal Policy Optimization (PPO) algorithm for reinforcement learning."""

from deepchem.models import TensorGraph
from deepchem.models.tensorgraph.optimizers import Adam
from deepchem.models.tensorgraph.layers import Feature, Weights, Label, Layer
import numpy as np
import tensorflow as tf
import collections
import copy
import multiprocessing
from multiprocessing.dummy import Pool
import os
import re
import time

class PPOLoss(Layer):
  """This layer computes the loss function for PPO."""

  def __init__(self, value_weight, entropy_weight, clipping_width, **kwargs):
    super(PPOLoss, self).__init__(**kwargs)
    self.value_weight = value_weight
    self.entropy_weight = entropy_weight
    self.clipping_width = clipping_width

  def create_tensor(self, **kwargs):
    reward, action, prob, value, advantage, old_prob = [
        layer.out_tensor for layer in self.in_layers
    ]
    machine_eps = np.finfo(np.float32).eps
    prob += machine_eps
    old_prob += machine_eps
    ratio = tf.reduce_sum(action * prob, axis=1) / old_prob
    clipped_ratio = tf.clip_by_value(ratio, 1 - self.clipping_width,
                                     1 + self.clipping_width)
    policy_loss = -tf.reduce_mean(
        tf.minimum(ratio * advantage, clipped_ratio * advantage))
    value_loss = tf.reduce_mean(tf.square(reward - value))
    entropy = -tf.reduce_mean(tf.reduce_sum(prob * tf.log(prob), axis=1))
    self.out_tensor = policy_loss + self.value_weight * value_loss - self.entropy_weight * entropy
    return self.out_tensor

class PPO(object):
  """
  Implements the Proximal Policy Optimization (PPO) algorithm for reinforcement learning.

  The algorithm is described in Schulman et al, "Proximal Policy Optimization Algorithms"
  (https://openai-public.s3-us-west-2.amazonaws.com/blog/2017-07/ppo/ppo-arxiv.pdf).
  This class requires the policy to output two quantities: a vector giving the probability of
  taking each action, and an estimate of the value function for the current state.  It
  optimizes both outputs at once using a loss that is the sum of three terms:

  1. The policy loss, which seeks to maximize the discounted reward for each action.
  2. The value loss, which tries to make the value estimate match the actual discounted reward
     that was attained at each step.
  3. An entropy term to encourage exploration.

  This class only supports environments with discrete action spaces, not continuous ones.  The
  "action" argument passed to the environment is an integer, giving the index of the action to perform.

  This class supports Generalized Advantage Estimation as described in Schulman et al., "High-Dimensional
  Continuous Control Using Generalized Advantage Estimation" (https://arxiv.org/abs/1506.02438).
  This is a method of trading off bias and variance in the advantage estimate, which can sometimes
  improve the rate of convergance.  Use the advantage_lambda parameter to adjust the tradeoff.

  This class supports Hindsight Experience Replay as described in Andrychowicz et al., "Hindsight
  Experience Replay" (https://arxiv.org/abs/1707.01495).  This is a method that can enormously
  accelerate learning when rewards are very rare.  It requires that the environment state contains
  information about the goal the agent is trying to achieve.  Each time it generates a rollout, it
  processes that rollout twice: once using the actual goal the agent was pursuing while generating
  it, and again using the final state of that rollout as the goal.  This guarantees that half of
  all rollouts processed will be ones that achieved their goals, and hence received a reward.

  To use this feature, specify use_hindsight=True to the constructor.  The environment must have
  a method defined as follows:

  def apply_hindsight(self, states, actions, goal):
    ...
    return new_states, rewards

  The method receives the list of states generated during the rollout, the action taken for each one,
  and a new goal state.  It should generate a new list of states that are identical to the input ones,
  except specifying the new goal.  It should return that list of states, and the rewards that would
  have been received for taking the specified actions from those states.
  """

  def __init__(self,
               env,
               policy,
               max_rollout_length=20,
               optimization_rollouts=8,
               optimization_epochs=4,
               batch_size=64,
               clipping_width=0.2,
               discount_factor=0.99,
               advantage_lambda=0.98,
               value_weight=1.0,
               entropy_weight=0.01,
               optimizer=None,
               model_dir=None,
               use_hindsight=False):
    """Create an object for optimizing a policy.

    Parameters
    ----------
    env: Environment
      the Environment to interact with
    policy: Policy
      the Policy to optimize.  Its create_layers() method must return a dict containing the
      keys 'action_prob' and 'value', corresponding to the action probabilities and value estimate
    max_rollout_length: int
      the maximum length of rollouts to generate
    optimization_rollouts: int
      the number of rollouts to generate for each iteration of optimization
    optimization_epochs: int
      the number of epochs of optimization to perform within each iteration
    batch_size: int
      the batch size to use during optimization.  If this is 0, each rollout will be used as a
      separate batch.
    clipping_width: float
      in computing the PPO loss function, the probability ratio is clipped to the range
      (1-clipping_width, 1+clipping_width)
    discount_factor: float
      the discount factor to use when computing rewards
    advantage_lambda: float
      the parameter for trading bias vs. variance in Generalized Advantage Estimation
    value_weight: float
      a scale factor for the value loss term in the loss function
    entropy_weight: float
      a scale factor for the entropy term in the loss function
    optimizer: Optimizer
      the optimizer to use.  If None, a default optimizer is used.
    model_dir: str
      the directory in which the model will be saved.  If None, a temporary directory will be created.
    use_hindsight: bool
      if True, use Hindsight Experience Replay
    """
    self._env = env
    self._policy = policy
    self.max_rollout_length = max_rollout_length
    self.optimization_rollouts = optimization_rollouts
    self.optimization_epochs = optimization_epochs
    self.batch_size = batch_size
    self.clipping_width = clipping_width
    self.discount_factor = discount_factor
    self.advantage_lambda = advantage_lambda
    self.value_weight = value_weight
    self.entropy_weight = entropy_weight
    self.use_hindsight = use_hindsight
    self._state_is_list = isinstance(env.state_shape[0], collections.Sequence)
    if optimizer is None:
      self._optimizer = Adam(learning_rate=0.001, beta1=0.9, beta2=0.999)
    else:
      self._optimizer = optimizer
    (self._graph, self._features, self._rewards, self._actions,
     self._action_prob, self._value, self._advantages,
     self._old_action_prob) = self._build_graph(None, 'global', model_dir)
    with self._graph._get_tf("Graph").as_default():
      self._session = tf.Session()
      self._train_op = self._graph._get_tf('Optimizer').minimize(
          self._graph.loss.out_tensor)
    self._rnn_states = self._graph.rnn_zero_states
    if len(self._rnn_states) > 0 and batch_size != 0:
      raise ValueError(
          'Cannot batch rollouts when the policy contains a recurrent layer.  Set batch_size to 0.'
      )

  def _build_graph(self, tf_graph, scope, model_dir):
    """Construct a TensorGraph containing the policy and loss calculations."""
    state_shape = self._env.state_shape
    state_dtype = self._env.state_dtype
    if not self._state_is_list:
      state_shape = [state_shape]
      state_dtype = [state_dtype]
    features = []
    for s, d in zip(state_shape, state_dtype):
      features.append(Feature(shape=[None] + list(s), dtype=tf.as_dtype(d)))
    policy_layers = self._policy.create_layers(features)
    action_prob = policy_layers['action_prob']
    value = policy_layers['value']
    rewards = Weights(shape=(None,))
    advantages = Weights(shape=(None,))
    old_action_prob = Weights(shape=(None,))
    actions = Label(shape=(None, self._env.n_actions))
    loss = PPOLoss(
        self.value_weight,
        self.entropy_weight,
        self.clipping_width,
        in_layers=[
            rewards, actions, action_prob, value, advantages, old_action_prob
        ])
    graph = TensorGraph(
        batch_size=self.max_rollout_length,
        use_queue=False,
        graph=tf_graph,
        model_dir=model_dir)
    for f in features:
      graph._add_layer(f)
    graph.add_output(action_prob)
    graph.add_output(value)
    graph.set_loss(loss)
    graph.set_optimizer(self._optimizer)
    with graph._get_tf("Graph").as_default():
      with tf.variable_scope(scope):
        graph.build()
    return graph, features, rewards, actions, action_prob, value, advantages, old_action_prob

  def fit(self,
          total_steps,
          max_checkpoints_to_keep=5,
          checkpoint_interval=600,
          restore=False):
    """Train the policy.

    Parameters
    ----------
    total_steps: int
      the total number of time steps to perform on the environment, across all rollouts
      on all threads
    max_checkpoints_to_keep: int
      the maximum number of checkpoint files to keep.  When this number is reached, older
      files are deleted.
    checkpoint_interval: float
      the time interval at which to save checkpoints, measured in seconds
    restore: bool
      if True, restore the model from the most recent checkpoint and continue training
      from there.  If False, retrain the model from scratch.
    """
    with self._graph._get_tf("Graph").as_default():
      step_count = 0
      workers = []
      threads = []
      for i in range(self.optimization_rollouts):
        workers.append(_Worker(self, i))
      self._session.run(tf.global_variables_initializer())
      if restore:
        self.restore()
      pool = Pool()
      variables = tf.get_collection(
          tf.GraphKeys.GLOBAL_VARIABLES, scope='global')
      saver = tf.train.Saver(variables, max_to_keep=max_checkpoints_to_keep)
      checkpoint_index = 0
      checkpoint_time = time.time()
      while step_count < total_steps:
        # Have the worker threads generate the rollouts for this iteration.

        rollouts = []
        pool.map(lambda x: rollouts.extend(x.run()), workers)

        # Perform optimization.

        for epoch in range(self.optimization_epochs):
          if self.batch_size == 0:
            batches = rollouts
          else:
            batches = self._iter_batches(rollouts)
          for batch in batches:
            initial_rnn_states, state_arrays, discounted_rewards, actions_matrix, action_prob, advantages = batch

            # Build the feed dict and run the optimizer.

            feed_dict = {}
            for placeholder, value in zip(self._graph.rnn_initial_states,
                                          initial_rnn_states):
              feed_dict[placeholder] = value
            for f, s in zip(self._features, state_arrays):
              feed_dict[f.out_tensor] = s
            feed_dict[self._rewards.out_tensor] = discounted_rewards
            feed_dict[self._actions.out_tensor] = actions_matrix
            feed_dict[self._advantages.out_tensor] = advantages
            feed_dict[self._old_action_prob.out_tensor] = action_prob
            feed_dict[self._graph.get_global_step()] = step_count
            self._session.run(self._train_op, feed_dict=feed_dict)

        # Update the number of steps taken so far and perform checkpointing.

        new_steps = sum(len(r[3]) for r in rollouts)
        if self.use_hindsight:
          new_steps /= 2
        step_count += new_steps
        if step_count >= total_steps or time.time(
        ) >= checkpoint_time + checkpoint_interval:
          saver.save(
              self._session,
              self._graph.save_file,
              global_step=checkpoint_index)
          checkpoint_index += 1
          checkpoint_time = time.time()

  def _iter_batches(self, rollouts):
    """Given a set of rollouts, merge them into batches for optimization."""

    # Merge all the rollouts into a single set of arrays.

    state_arrays = []
    for i in range(len(rollouts[0][1])):
      state_arrays.append(np.concatenate([r[1][i] for r in rollouts]))
    discounted_rewards = np.concatenate([r[2] for r in rollouts])
    actions_matrix = np.concatenate([r[3] for r in rollouts])
    action_prob = np.concatenate([r[4] for r in rollouts])
    advantages = np.concatenate([r[5] for r in rollouts])
    total_length = len(discounted_rewards)

    # Iterate slices.

    start = 0
    while start < total_length:
      end = min(start + self.batch_size, total_length)
      batch = [[]]
      batch.append([s[start:end] for s in state_arrays])
      batch.append(discounted_rewards[start:end])
      batch.append(actions_matrix[start:end])
      batch.append(action_prob[start:end])
      batch.append(advantages[start:end])
      start = end
      yield batch

  def predict(self, state, use_saved_states=True, save_states=True):
    """Compute the policy's output predictions for a state.

    If the policy involves recurrent layers, this method can preserve their internal
    states between calls.  Use the use_saved_states and save_states arguments to specify
    how it should behave.

    Parameters
    ----------
    state: array
      the state of the environment for which to generate predictions
    use_saved_states: bool
      if True, the states most recently saved by a previous call to predict() or select_action()
      will be used as the initial states.  If False, the internal states of all recurrent layers
      will be set to all zeros before computing the predictions.
    save_states: bool
      if True, the internal states of all recurrent layers at the end of the calculation
      will be saved, and any previously saved states will be discarded.  If False, the
      states at the end of the calculation will be discarded, and any previously saved
      states will be kept.

    Returns
    -------
    the array of action probabilities, and the estimated value function
    """
    if not self._state_is_list:
      state = [state]
    with self._graph._get_tf("Graph").as_default():
      feed_dict = self._create_feed_dict(state, use_saved_states)
      tensors = [self._action_prob.out_tensor, self._value.out_tensor]
      if save_states:
        tensors += self._graph.rnn_final_states
      results = self._session.run(tensors, feed_dict=feed_dict)
      if save_states:
        self._rnn_states = results[2:]
      return results[:2]

  def select_action(self,
                    state,
                    deterministic=False,
                    use_saved_states=True,
                    save_states=True):
    """Select an action to perform based on the environment's state.

    If the policy involves recurrent layers, this method can preserve their internal
    states between calls.  Use the use_saved_states and save_states arguments to specify
    how it should behave.

    Parameters
    ----------
    state: array
      the state of the environment for which to select an action
    deterministic: bool
      if True, always return the best action (that is, the one with highest probability).
      If False, randomly select an action based on the computed probabilities.
    use_saved_states: bool
      if True, the states most recently saved by a previous call to predict() or select_action()
      will be used as the initial states.  If False, the internal states of all recurrent layers
      will be set to all zeros before computing the predictions.
    save_states: bool
      if True, the internal states of all recurrent layers at the end of the calculation
      will be saved, and any previously saved states will be discarded.  If False, the
      states at the end of the calculation will be discarded, and any previously saved
      states will be kept.

    Returns
    -------
    the index of the selected action
    """
    if not self._state_is_list:
      state = [state]
    with self._graph._get_tf("Graph").as_default():
      feed_dict = self._create_feed_dict(state, use_saved_states)
      tensors = [self._action_prob.out_tensor]
      if save_states:
        tensors += self._graph.rnn_final_states
      results = self._session.run(tensors, feed_dict=feed_dict)
      probabilities = results[0]
      if save_states:
        self._rnn_states = results[1:]
      if deterministic:
        return probabilities.argmax()
      else:
        return np.random.choice(
            np.arange(self._env.n_actions), p=probabilities[0])

  def restore(self):
    """Reload the model parameters from the most recent checkpoint file."""
    last_checkpoint = tf.train.latest_checkpoint(self._graph.model_dir)
    if last_checkpoint is None:
      raise ValueError('No checkpoint found')
    with self._graph._get_tf("Graph").as_default():
      variables = tf.get_collection(
          tf.GraphKeys.GLOBAL_VARIABLES, scope='global')
      saver = tf.train.Saver(variables)
      saver.restore(self._session, last_checkpoint)

  def _create_feed_dict(self, state, use_saved_states):
    """Create a feed dict for use by predict() or select_action()."""
    feed_dict = dict((f.out_tensor, np.expand_dims(s, axis=0))
                     for f, s in zip(self._features, state))
    if use_saved_states:
      rnn_states = self._rnn_states
    else:
      rnn_states = self._graph.rnn_zero_states
    for (placeholder, value) in zip(self._graph.rnn_initial_states, rnn_states):
      feed_dict[placeholder] = value
    return feed_dict

class _Worker(object):
  """A Worker object is created for each training thread."""

  def __init__(self, ppo, index):
    self.ppo = ppo
    self.index = index
    self.scope = 'worker%d' % index
    self.env = copy.deepcopy(ppo._env)
    self.env.reset()
    self.graph, self.features, self.rewards, self.actions, self.action_prob, self.value, self.advantages, self.old_action_prob = ppo._build_graph(
        ppo._graph._get_tf('Graph'), self.scope, None)
    self.rnn_states = self.graph.rnn_zero_states
    with ppo._graph._get_tf("Graph").as_default():
      local_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
                                     self.scope)
      global_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
                                      'global')
      self.update_local_variables = tf.group(
          *[tf.assign(v1, v2) for v1, v2 in zip(local_vars, global_vars)])

  def run(self):
    rollouts = []
    with self.graph._get_tf("Graph").as_default():
      self.ppo._session.run(self.update_local_variables)
      initial_rnn_states = self.rnn_states
      states, actions, action_prob, rewards, values = self.create_rollout()
      rollouts.append(
          self.process_rollout(states, actions, action_prob, rewards, values,
                               initial_rnn_states))
      if self.ppo.use_hindsight:
        rollouts.append(
            self.process_rollout_with_hindsight(states, actions,
                                                initial_rnn_states))
    return rollouts

  def create_rollout(self):
    """Generate a rollout."""
    n_actions = self.env.n_actions
    session = self.ppo._session
    states = []
    action_prob = []
    actions = []
    rewards = []
    values = []

    # Generate the rollout.

    for i in range(self.ppo.max_rollout_length):
      if self.env.terminated:
        break
      state = self.env.state
      states.append(state)
      feed_dict = self.create_feed_dict(state)
      results = session.run(
          [self.action_prob.out_tensor, self.value.out_tensor] +
          self.graph.rnn_final_states,
          feed_dict=feed_dict)
      probabilities, value = results[:2]

      self.rnn_states = results[2:]
      action = np.random.choice(np.arange(n_actions), p=probabilities[0])
      actions.append(action)
      action_prob.append(probabilities[0][action])
      values.append(float(value))
      rewards.append(self.env.step(action))

    # Compute an estimate of the reward for the rest of the episode.

    if not self.env.terminated:
      feed_dict = self.create_feed_dict(self.env.state)
      final_value = self.ppo.discount_factor * float(
          session.run(self.value.out_tensor, feed_dict))
    else:
      final_value = 0.0
    values.append(final_value)
    if self.env.terminated:
      self.env.reset()
      self.rnn_states = self.graph.rnn_zero_states
    return states, np.array(
        actions, dtype=np.int32), np.array(action_prob), np.array(
            rewards), np.array(values)

  def process_rollout(self, states, actions, action_prob, rewards, values,
                      initial_rnn_states):
    """Construct the arrays needed for training."""

    # Compute the discounted rewards and advantages.

    discounted_rewards = rewards.copy()
    discounted_rewards[-1] += values[-1]
    advantages = rewards - values[:-1] + self.ppo.discount_factor * np.array(
        values[1:])
    for j in range(len(rewards) - 1, 0, -1):
      discounted_rewards[j -
                         1] += self.ppo.discount_factor * discounted_rewards[j]
      advantages[
          j -
          1] += self.ppo.discount_factor * self.ppo.advantage_lambda * advantages[
              j]

    # Convert the actions to one-hot.

    n_actions = self.env.n_actions
    actions_matrix = []
    for action in actions:
      a = np.zeros(n_actions)
      a[action] = 1.0
      actions_matrix.append(a)

    # Rearrange the states into the proper set of arrays.

    if self.ppo._state_is_list:
      state_arrays = [[] for i in range(len(self.features))]
      for state in states:
        for j in range(len(state)):
          state_arrays[j].append(state[j])
    else:
      state_arrays = [states]

    # Return the processed arrays.

    return (initial_rnn_states, state_arrays, discounted_rewards,
            actions_matrix, action_prob, advantages)

  def process_rollout_with_hindsight(self, states, actions, initial_rnn_states):
    """Create a new rollout by applying hindsight to an existing one, then process it."""
    hindsight_states, rewards = self.env.apply_hindsight(
        states, actions, states[-1])
    if self.ppo._state_is_list:
      state_arrays = [[] for i in range(len(self.features))]
      for state in hindsight_states:
        for j in range(len(state)):
          state_arrays[j].append(state[j])
    else:
      state_arrays = [hindsight_states]
    feed_dict = {}
    for placeholder, value in zip(self.graph.rnn_initial_states,
                                  initial_rnn_states):
      feed_dict[placeholder] = value
    for f, s in zip(self.features, state_arrays):
      feed_dict[f.out_tensor] = s
    values, probabilities = self.ppo._session.run(
        [self.value.out_tensor, self.action_prob.out_tensor],
        feed_dict=feed_dict)
    values = np.append(values.flatten(), 0.0)
    action_prob = probabilities[np.arange(len(actions)), actions]
    return self.process_rollout(hindsight_states, actions, action_prob,
                                np.array(rewards),
                                np.array(values), initial_rnn_states)

  def create_feed_dict(self, state):
    """Create a feed dict for use during a rollout."""
    if not self.ppo._state_is_list:
      state = [state]
    feed_dict = dict((f.out_tensor, np.expand_dims(s, axis=0))
                     for f, s in zip(self.features, state))
    for (placeholder, value) in zip(self.graph.rnn_initial_states,
                                    self.rnn_states):
      feed_dict[placeholder] = value
    return feed_dict


text:
import numpy as np
from sklearn.decomposition import SparsePCA

def sparse_pca(K, alpha, ridge_alpha):
    transformer = SparsePCA(n_components=1, alpha=alpha, ridge_alpha=ridge_alpha, normalize_components=False, random_state=0)
    transformer.fit(K)
    val = transformer.components_[0]
    print('#nnz: ', np.sum(np.abs(val) > 1.0e-10))
    #print(np.sum(val * val))
    #val = np.random.randn(K.shape[1])
    return val / np.linalg.norm(val)

ext: py

reason: import numpy


text:
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from gekko import GEKKO

# Import or generate data
filename = 'tclab_dyn_data2.csv'
try:
    data = pd.read_csv(filename)
except:
    url = 'https://apmonitor.com/do/uploads/Main/tclab_dyn_data2.txt'
    data = pd.read_csv(url)

# Create GEKKO Model
m = GEKKO()
m.time = data['Time'].values

# Parameters to Estimate
U = m.FV(value=10,lb=1,ub=20)
tau = m.FV(value=20,lb=15,ub=25)
alpha1 = m.FV(value=0.01,lb=0.003,ub=0.03)  # W / % heater
alpha2 = m.FV(value=0.005,lb=0.002,ub=0.02) # W / % heater

# STATUS=1 allows solver to adjust parameter
U.STATUS = 1
tau.STATUS = 1
alpha1.STATUS = 1
alpha2.STATUS = 1

# Measured inputs
Q1 = m.MV(value=data['H1'].values)
Q2 = m.MV(value=data['H2'].values)

# State variables
TH1 = m.SV(value=data['T1'].values)
TH2 = m.SV(value=data['T2'].values)

# Measurements for model alignment
TC1 = m.CV(value=data['T1'].values,lb=0,ub=200)
TC1.FSTATUS = 1    # minimize fstatus * (meas-pred)^2

TC2 = m.CV(value=data['T2'].values,lb=0,ub=200)
TC2.FSTATUS = 1    # minimize fstatus * (meas-pred)^2

Ta = m.Param(value=19.0+273.15)     # K
mass = m.Param(value=4.0/1000.0)    # kg
Cp = m.Param(value=0.5*1000.0)      # J/kg-K
A = m.Param(value=10.0/100.0**2)    # Area not between heaters in m^2
As = m.Param(value=2.0/100.0**2)    # Area between heaters in m^2
eps = m.Param(value=0.9)            # Emissivity
sigma = m.Const(5.67e-8)            # Stefan-Boltzmann

# Heater temperatures in Kelvin
T1 = m.Intermediate(TH1+273.15)
T2 = m.Intermediate(TH2+273.15)

# Heat transfer between two heaters
Q_C12 = m.Intermediate(U*As*(T2-T1)) # Convective
Q_R12 = m.Intermediate(eps*sigma*As*(T2**4-T1**4)) # Radiative

# Semi-fundamental correlations (energy balances)
m.Equation(TH1.dt() == (1.0/(mass*Cp))*(U*A*(Ta-T1) \
                    + eps * sigma * A * (Ta**4 - T1**4) \
                    + Q_C12 + Q_R12 \
                    + alpha1*Q1))

m.Equation(TH2.dt() == (1.0/(mass*Cp))*(U*A*(Ta-T2) \
                    + eps * sigma * A * (Ta**4 - T2**4) \
                    - Q_C12 - Q_R12 \
                    + alpha2*Q2))

# Empirical correlations (lag equations to emulate conduction)
m.Equation(tau * TC1.dt() == -TC1 + TH1)
m.Equation(tau * TC2.dt() == -TC2 + TH2)

# Options
m.options.IMODE   = 5 # MHE
m.options.EV_TYPE = 2 # Objective type
m.options.NODES   = 2 # Collocation nodes
m.options.SOLVER  = 3 # IPOPT

# Solve
m.solve(disp=True)

# Parameter values
print('U     : ' + str(U.value[0]))
print('tau   : ' + str(tau.value[0]))
print('alpha1: ' + str(alpha1.value[0]))
print('alpha2: ' + str(alpha2.value[0]))

# Create plot
plt.figure()
ax=plt.subplot(2,1,1)
ax.grid()
plt.plot(data['Time'],data['T1'],'ro',label=r'$T_1$ measured')
plt.plot(m.time,TH1.value,'k-',label=r'$T_1$ heater')
plt.plot(m.time,TC1.value,color='purple',linestyle='--',\
         linewidth=3,label=r'$T_1$ sensor')
plt.plot(data['Time'],data['T2'],'bx',label=r'$T_2$ measured')
plt.plot(m.time,TH2.value,'k-',label=r'$T_2$ heater')
plt.plot(m.time,TC2.value,color='orange',linestyle='--',\
         linewidth=3,label=r'$T_2$ sensor')
plt.ylabel('Temperature (degC)')
plt.legend(loc=2)
ax=plt.subplot(2,1,2)
ax.grid()
plt.plot(data['Time'],data['H1'],'r-',\
         linewidth=3,label=r'$Q_1$')
plt.plot(data['Time'],data['H2'],'b:',\
         linewidth=3,label=r'$Q_2$')
plt.ylabel('Heaters')
plt.xlabel('Time (sec)')
plt.legend(loc='best')
plt.show()

max_stars_repo_name: APMonitor/arduino

max_issues_repo_name: rezabbasi/arduino

max_forks_repo_name: rezabbasi/arduino


text:
# This code is part of Qiskit.
#
# (C) Copyright IBM 2017, 2019.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

"""
DensityMatrix quantum state class.
"""

import warnings
from numbers import Number
import numpy as np

from qiskit.circuit.quantumcircuit import QuantumCircuit
from qiskit.circuit.instruction import Instruction
from qiskit.exceptions import QiskitError
from qiskit.quantum_info.states.quantum_state import QuantumState
from qiskit.quantum_info.operators.tolerances import TolerancesMixin
from qiskit.quantum_info.operators.operator import Operator
from qiskit.quantum_info.operators.scalar_op import ScalarOp
from qiskit.quantum_info.operators.predicates import is_hermitian_matrix
from qiskit.quantum_info.operators.predicates import is_positive_semidefinite_matrix
from qiskit.quantum_info.operators.channel.quantum_channel import QuantumChannel
from qiskit.quantum_info.operators.channel.superop import SuperOp
from qiskit.quantum_info.states.statevector import Statevector

class DensityMatrix(QuantumState, TolerancesMixin):
    """DensityMatrix class"""

    def __init__(self, data, dims=None):
        """Initialize a density matrix object.

        Args:
            data (matrix_like or vector_like): a density matrix or
                statevector. If a vector the density matrix is constructed
                as the projector of that vector.
            dims (int or tuple or list): Optional. The subsystem dimension
                    of the state (See additional information).

        Raises:
            QiskitError: if input data is not valid.

        Additional Information:
            The ``dims`` kwarg can be None, an integer, or an iterable of
            integers.

            * ``Iterable`` -- the subsystem dimensions are the values in the list
              with the total number of subsystems given by the length of the list.

            * ``Int`` or ``None`` -- the leading dimension of the input matrix
              specifies the total dimension of the density matrix. If it is a
              power of two the state will be initialized as an N-qubit state.
              If it is not a power of two the state will have a single
              d-dimensional subsystem.
        """
        if isinstance(data, (list, np.ndarray)):
            # Finally we check if the input is a raw matrix in either a
            # python list or numpy array format.
            self._data = np.asarray(data, dtype=complex)
        elif hasattr(data, 'to_operator'):
            # If the data object has a 'to_operator' attribute this is given
            # higher preference than the 'to_matrix' method for initializing
            # an Operator object.
            op = data.to_operator()
            self._data = op.data
            if dims is None:
                dims = op._output_dims
        elif hasattr(data, 'to_matrix'):
            # If no 'to_operator' attribute exists we next look for a
            # 'to_matrix' attribute to a matrix that will be cast into
            # a complex numpy matrix.
            self._data = np.asarray(data.to_matrix(), dtype=complex)
        else:
            raise QiskitError("Invalid input data format for DensityMatrix")
        # Convert statevector into a density matrix
        ndim = self._data.ndim
        shape = self._data.shape
        if ndim == 2 and shape[0] == shape[1]:
            pass  # We good
        elif ndim == 1:
            self._data = np.outer(self._data, np.conj(self._data))
        elif ndim == 2 and shape[1] == 1:
            self._data = np.reshape(self._data, shape[0])
            shape = self._data.shape
        else:
            raise QiskitError(
                "Invalid DensityMatrix input: not a square matrix.")
        super().__init__(self._automatic_dims(dims, shape[0]))

    def __eq__(self, other):
        return super().__eq__(other) and np.allclose(
            self._data, other._data, rtol=self.rtol, atol=self.atol)

    def __repr__(self):
        prefix = 'DensityMatrix('
        pad = len(prefix) * ' '
        return '{}{},\n{}dims={})'.format(
            prefix, np.array2string(
                self._data, separator=', ', prefix=prefix),
            pad, self._dims)

    @property
    def data(self):
        """Return data."""
        return self._data

    def is_valid(self, atol=None, rtol=None):
        """Return True if trace 1 and positive semidefinite."""
        if atol is None:
            atol = self.atol
        if rtol is None:
            rtol = self.rtol
        # Check trace == 1
        if not np.allclose(self.trace(), 1, rtol=rtol, atol=atol):
            return False
        # Check Hermitian
        if not is_hermitian_matrix(self.data, rtol=rtol, atol=atol):
            return False
        # Check positive semidefinite
        return is_positive_semidefinite_matrix(self.data, rtol=rtol, atol=atol)

    def to_operator(self):
        """Convert to Operator"""
        dims = self.dims()
        return Operator(self.data, input_dims=dims, output_dims=dims)

    def conjugate(self):
        """Return the conjugate of the density matrix."""
        return DensityMatrix(np.conj(self.data), dims=self.dims())

    def trace(self):
        """Return the trace of the density matrix."""
        return np.trace(self.data)

    def purity(self):
        """Return the purity of the quantum state."""
        # For a valid statevector the purity is always 1, however if we simply
        # have an arbitrary vector (not correctly normalized) then the
        # purity is equivalent to the trace squared:
        # P(|psi>) = Tr[|psi><psi|psi><psi|] = |<psi|psi>|^2
        return np.trace(np.dot(self.data, self.data))

    def tensor(self, other):
        """Return the tensor product state self ⊗ other.

        Args:
            other (DensityMatrix): a quantum state object.

        Returns:
            DensityMatrix: the tensor product operator self ⊗ other.

        Raises:
            QiskitError: if other is not a quantum state.
        """
        if not isinstance(other, DensityMatrix):
            other = DensityMatrix(other)
        dims = other.dims() + self.dims()
        data = np.kron(self._data, other._data)
        return DensityMatrix(data, dims)

    def expand(self, other):
        """Return the tensor product state other ⊗ self.

        Args:
            other (DensityMatrix): a quantum state object.

        Returns:
            DensityMatrix: the tensor product state other ⊗ self.

        Raises:
            QiskitError: if other is not a quantum state.
        """
        if not isinstance(other, DensityMatrix):
            other = DensityMatrix(other)
        dims = self.dims() + other.dims()
        data = np.kron(other._data, self._data)
        return DensityMatrix(data, dims)

    def _add(self, other):
        """Return the linear combination self + other.

        Args:
            other (DensityMatrix): a quantum state object.

        Returns:
            DensityMatrix: the linear combination self + other.

        Raises:
            QiskitError: if other is not a quantum state, or has
                         incompatible dimensions.
        """
        if not isinstance(other, DensityMatrix):
            other = DensityMatrix(other)
        if self.dim != other.dim:
            raise QiskitError("other DensityMatrix has different dimensions.")
        return DensityMatrix(self.data + other.data, self.dims())

    def _multiply(self, other):
        """Return the scalar multiplied state other * self.

        Args:
            other (complex): a complex number.

        Returns:
            DensityMatrix: the scalar multiplied state other * self.

        Raises:
            QiskitError: if other is not a valid complex number.
        """
        if not isinstance(other, Number):
            raise QiskitError("other is not a number")
        return DensityMatrix(other * self.data, self.dims())

    def evolve(self, other, qargs=None):
        """Evolve a quantum state by an operator.

        Args:
            other (Operator or QuantumChannel
                   or Instruction or Circuit): The operator to evolve by.
            qargs (list): a list of QuantumState subsystem positions to apply
                           the operator on.

        Returns:
            QuantumState: the output quantum state.

        Raises:
            QiskitError: if the operator dimension does not match the
                         specified QuantumState subsystem dimensions.
        """
        if qargs is None:
            qargs = getattr(other, 'qargs', None)

        # Evolution by a circuit or instruction
        if isinstance(other, (QuantumCircuit, Instruction)):
            return self._evolve_instruction(other, qargs=qargs)

        # Evolution by a QuantumChannel
        if hasattr(other, 'to_quantumchannel'):
            return other.to_quantumchannel()._evolve(self, qargs=qargs)
        if isinstance(other, QuantumChannel):
            return other._evolve(self, qargs=qargs)

        # Unitary evolution by an Operator
        if not isinstance(other, Operator):
            other = Operator(other)
        return self._evolve_operator(other, qargs=qargs)

    def expectation_value(self, oper, qargs=None):
        """Compute the expectation value of an operator.

        Args:
            oper (Operator): an operator to evaluate expval.
            qargs (None or list): subsystems to apply the operator on.

        Returns:
            complex: the expectation value.
        """
        if not isinstance(oper, Operator):
            oper = Operator(oper)
        return np.trace(Operator(self).dot(oper.adjoint(), qargs=qargs).data)

    def probabilities(self, qargs=None, decimals=None):
        """Return the subsystem measurement probability vector.

        Measurement probabilities are with respect to measurement in the
        computation (diagonal) basis.

        Args:
            qargs (None or list): subsystems to return probabilities for,
                if None return for all subsystems (Default: None).
            decimals (None or int): the number of decimal places to round
                values. If None no rounding is done (Default: None).

        Returns:
            np.array: The Numpy vector array of probabilities.

        Examples:

            Consider a 2-qubit product state :math:`\\rho=\\rho_1\\otimes\\rho_0`
            with :math:`\\rho_1=|+\\rangle\\!\\langle+|`,
            :math:`\\rho_0=|0\\rangle\\!\\langle0|`.

            .. jupyter-execute::

                from qiskit.quantum_info import DensityMatrix

                rho = DensityMatrix.from_label('+0')

                # Probabilities for measuring both qubits
                probs = rho.probabilities()
                print('probs: {}'.format(probs))

                # Probabilities for measuring only qubit-0
                probs_qubit_0 = rho.probabilities([0])
                print('Qubit-0 probs: {}'.format(probs_qubit_0))

                # Probabilities for measuring only qubit-1
                probs_qubit_1 = rho.probabilities([1])
                print('Qubit-1 probs: {}'.format(probs_qubit_1))

            We can also permute the order of qubits in the ``qargs`` list
            to change the qubit position in the probabilities output

            .. jupyter-execute::

                from qiskit.quantum_info import DensityMatrix

                rho = DensityMatrix.from_label('+0')

                # Probabilities for measuring both qubits
                probs = rho.probabilities([0, 1])
                print('probs: {}'.format(probs))

                # Probabilities for measuring both qubits
                # but swapping qubits 0 and 1 in output
                probs_swapped = rho.probabilities([1, 0])
                print('Swapped probs: {}'.format(probs_swapped))
        """
        probs = self._subsystem_probabilities(
            np.abs(self.data.diagonal()), self._dims, qargs=qargs)
        if decimals is not None:
            probs = probs.round(decimals=decimals)
        return probs

    def reset(self, qargs=None):
        """Reset state or subsystems to the 0-state.

        Args:
            qargs (list or None): subsystems to reset, if None all
                                  subsystems will be reset to their 0-state
                                  (Default: None).

        Returns:
            DensityMatrix: the reset state.

        Additional Information:
            If all subsystems are reset this will return the ground state
            on all subsystems. If only a some subsystems are reset this
            function will perform evolution by the reset
            :class:`~qiskit.quantum_info.SuperOp` of the reset subsystems.
        """
        if qargs is None:
            # Resetting all qubits does not require sampling or RNG
            state = np.zeros(2 * (self._dim, ), dtype=complex)
            state[0, 0] = 1
            return DensityMatrix(state, dims=self._dims)

        # Reset by evolving by reset SuperOp
        dims = self.dims(qargs)
        reset_superop = SuperOp(ScalarOp(dims, coeff=0))
        reset_superop.data[0] = Operator(ScalarOp(dims)).data.ravel()
        return self.evolve(reset_superop, qargs=qargs)

    @classmethod
    def from_label(cls, label):
        r"""Return a tensor product of Pauli X,Y,Z eigenstates.

        .. list-table:: Single-qubit state labels
           :header-rows: 1

           * - Label
             - Statevector
           * - ``"0"``
             - :math:`\begin{pmatrix} 1 & 0 \\ 0 & 0 \end{pmatrix}`
           * - ``"1"``
             - :math:`\begin{pmatrix} 0 & 0 \\ 0 & 1 \end{pmatrix}`
           * - ``"+"``
             - :math:`\frac{1}{2}\begin{pmatrix} 1 & 1 \\ 1 & 1 \end{pmatrix}`
           * - ``"-"``
             - :math:`\frac{1}{2}\begin{pmatrix} 1 & -1 \\ -1 & 1 \end{pmatrix}`
           * - ``"r"``
             - :math:`\frac{1}{2}\begin{pmatrix} 1 & -i \\ i & 1 \end{pmatrix}`
           * - ``"l"``
             - :math:`\frac{1}{2}\begin{pmatrix} 1 & i \\ -i & 1 \end{pmatrix}`

        Args:
            label (string): a eigenstate string ket label (see table for
                            allowed values).

        Returns:
            Statevector: The N-qubit basis state density matrix.

        Raises:
            QiskitError: if the label contains invalid characters, or the length
                         of the label is larger than an explicitly specified num_qubits.
        """
        return DensityMatrix(Statevector.from_label(label))

    @staticmethod
    def from_int(i, dims):
        """Return a computational basis state density matrix.

        Args:
            i (int): the basis state element.
            dims (int or tuple or list): The subsystem dimensions of the statevector
                                         (See additional information).

        Returns:
            DensityMatrix: The computational basis state :math:`|i\\rangle\\!\\langle i|`.

        Additional Information:
            The ``dims`` kwarg can be an integer or an iterable of integers.

            * ``Iterable`` -- the subsystem dimensions are the values in the list
              with the total number of subsystems given by the length of the list.

            * ``Int`` -- the integer specifies the total dimension of the
              state. If it is a power of two the state will be initialized
              as an N-qubit state. If it is not a power of  two the state
              will have a single d-dimensional subsystem.
        """
        size = np.product(dims)
        state = np.zeros((size, size), dtype=complex)
        state[i, i] = 1.0
        return DensityMatrix(state, dims=dims)

    @classmethod
    def from_instruction(cls, instruction):
        """Return the output density matrix of an instruction.

        The statevector is initialized in the state :math:`|{0,\\ldots,0}\\rangle` of
        the same number of qubits as the input instruction or circuit, evolved
        by the input instruction, and the output statevector returned.

        Args:
            instruction (qiskit.circuit.Instruction or QuantumCircuit): instruction or circuit

        Returns:
            DensityMatrix: the final density matrix.

        Raises:
            QiskitError: if the instruction contains invalid instructions for
                         density matrix simulation.
        """
        # Convert circuit to an instruction
        if isinstance(instruction, QuantumCircuit):
            instruction = instruction.to_instruction()
        # Initialize an the statevector in the all |0> state
        num_qubits = instruction.num_qubits
        init = np.zeros((2**num_qubits, 2**num_qubits), dtype=complex)
        init[0, 0] = 1
        vec = DensityMatrix(init, dims=num_qubits * (2, ))
        vec._append_instruction(instruction)
        return vec

    def to_dict(self, decimals=None):
        r"""Convert the density matrix to dictionary form.

        This dictionary representation uses a Ket-like notation where the
        dictionary keys are qudit strings for the subsystem basis vectors.
        If any subsystem has a dimension greater than 10 comma delimiters are
        inserted between integers so that subsystems can be distinguished.

        Args:
            decimals (None or int): the number of decimal places to round
                                    values. If None no rounding is done
                                    (Default: None).

        Returns:
            dict: the dictionary form of the DensityMatrix.

        Examples:

            The ket-form of a 2-qubit density matrix
            :math:`rho = |-\rangle\!\langle -|\otimes |0\rangle\!\langle 0|`

            .. jupyter-execute::

                from qiskit.quantum_info import DensityMatrix

                rho = DensityMatrix.from_label('-0')
                print(rho.to_dict())

            For non-qubit subsystems the integer range can go from 0 to 9. For
            example in a qutrit system

            .. jupyter-execute::

                import numpy as np
                from qiskit.quantum_info import DensityMatrix

                mat = np.zeros((9, 9))
                mat[0, 0] = 0.25
                mat[3, 3] = 0.25
                mat[6, 6] = 0.25
                mat[-1, -1] = 0.25
                rho = DensityMatrix(mat, dims=(3, 3))
                print(rho.to_dict())

            For large subsystem dimensions delimeters are required. The
            following example is for a 20-dimensional system consisting of
            a qubit and 10-dimensional qudit.

            .. jupyter-execute::

                import numpy as np
                from qiskit.quantum_info import DensityMatrix

                mat = np.zeros((2 * 10, 2 * 10))
                mat[0, 0] = 0.5
                mat[-1, -1] = 0.5
                rho = DensityMatrix(mat, dims=(2, 10))
                print(rho.to_dict())
        """
        return self._matrix_to_dict(self.data,
                                    self._dims,
                                    decimals=decimals,
                                    string_labels=True)

    @property
    def _shape(self):
        """Return the tensor shape of the matrix operator"""
        return 2 * tuple(reversed(self.dims()))

    def _evolve_operator(self, other, qargs=None):
        """Evolve density matrix by an operator"""
        if qargs is None:
            # Evolution on full matrix
            if self._dim != other.dim[0]:
                raise QiskitError(
                    "Operator input dimension is not equal to density matrix dimension."
                )
            op_mat = other.data
            mat = np.dot(op_mat, self.data).dot(op_mat.T.conj())
            return DensityMatrix(mat, dims=other.output_dims())
        # Otherwise we are applying an operator only to subsystems
        # Check dimensions of subsystems match the operator
        if self.dims(qargs) != other.input_dims():
            raise QiskitError(
                "Operator input dimensions are not equal to statevector subsystem dimensions."
            )
        # Reshape statevector and operator
        tensor = np.reshape(self.data, self._shape)
        # Construct list of tensor indices of statevector to be contracted
        num_indices = len(self.dims())
        indices = [num_indices - 1 - qubit for qubit in qargs]
        # Left multiple by mat
        mat = np.reshape(other.data, other._shape)
        tensor = Operator._einsum_matmul(tensor, mat, indices)
        # Right multiply by mat ** dagger
        adj = other.adjoint()
        mat_adj = np.reshape(adj.data, adj._shape)
        tensor = Operator._einsum_matmul(tensor, mat_adj, indices, num_indices,
                                         True)
        # Replace evolved dimensions
        new_dims = list(self.dims())
        for qubit, dim in zip(qargs, other.output_dims()):
            new_dims[qubit] = dim
        new_dim = np.product(new_dims)
        return DensityMatrix(np.reshape(tensor, (new_dim, new_dim)),
                             dims=new_dims)

    def _append_instruction(self, other, qargs=None):
        """Update the current Statevector by applying an instruction."""
        from qiskit.circuit.reset import Reset
        from qiskit.circuit.barrier import Barrier

        # Try evolving by a matrix operator (unitary-like evolution)
        mat = Operator._instruction_to_matrix(other)
        if mat is not None:
            self._data = self._evolve_operator(Operator(mat), qargs=qargs).data
            return

        # Special instruction types
        if isinstance(other, Reset):
            self._data = self.reset(qargs)._data
            return
        if isinstance(other, Barrier):
            return

        # Otherwise try evolving by a Superoperator
        chan = SuperOp._instruction_to_superop(other)
        if chan is not None:
            # Evolve current state by the superoperator
            self._data = chan._evolve(self, qargs=qargs).data
            return
        # If the instruction doesn't have a matrix defined we use its
        # circuit decomposition definition if it exists, otherwise we
        # cannot compose this gate and raise an error.
        if other.definition is None:
            raise QiskitError('Cannot apply Instruction: {}'.format(
                other.name))
        if not isinstance(other.definition, QuantumCircuit):
            raise QiskitError('{} instruction definition is {}; expected QuantumCircuit'.format(
                other.name, type(other.definition)))
        for instr, qregs, cregs in other.definition:
            if cregs:
                raise QiskitError(
                    'Cannot apply instruction with classical registers: {}'.
                    format(instr.name))
            # Get the integer position of the flat register
            if qargs is None:
                new_qargs = [tup.index for tup in qregs]
            else:
                new_qargs = [qargs[tup.index] for tup in qregs]
            self._append_instruction(instr, qargs=new_qargs)

    def _evolve_instruction(self, obj, qargs=None):
        """Return a new statevector by applying an instruction."""
        if isinstance(obj, QuantumCircuit):
            obj = obj.to_instruction()
        vec = DensityMatrix(self.data, dims=self._dims)
        vec._append_instruction(obj, qargs=qargs)
        return vec

    def to_statevector(self, atol=None, rtol=None):
        """Return a statevector from a pure density matrix.

        Args:
            atol (float): Absolute tolerance for checking operation validity.
            rtol (float): Relative tolerance for checking operation validity.

        Returns:
            Statevector: The pure density matrix's corresponding statevector.
                Corresponds to the eigenvector of the only non-zero eigenvalue.

        Raises:
            QiskitError: if the state is not pure.
        """
        if atol is None:
            atol = self.atol
        if rtol is None:
            rtol = self.rtol

        if not is_hermitian_matrix(self._data, atol=atol, rtol=rtol):
            raise QiskitError("Not a valid density matrix (non-hermitian).")

        evals, evecs = np.linalg.eig(self._data)

        nonzero_evals = evals[abs(evals) > atol]
        if len(nonzero_evals) != 1 or not np.isclose(nonzero_evals[0], 1,
                                                     atol=atol, rtol=rtol):
            raise QiskitError("Density matrix is not a pure state")

        psi = evecs[:, np.argmax(evals)]  # eigenvectors returned in columns.
        return Statevector(psi)

    def to_counts(self):
        """Returns the density matrix as a counts dict of probabilities.

        DEPRECATED: use :meth:`probabilities_dict` instead.

        Returns:
            dict: Counts of probabilities.
        """
        warnings.warn(
            'The `Statevector.to_counts` method is deprecated as of 0.13.0,'
            ' and will be removed no earlier than 3 months after that '
            'release date. You should use the `Statevector.probabilities_dict`'
            ' method instead.', DeprecationWarning, stacklevel=2)
        return self.probabilities_dict()

max_stars_repo_name: septembrr/qiskit-terra

max_issues_repo_name: septembrr/qiskit-terra

max_forks_repo_name: septembrr/qiskit-terra


text:
# select_random_subset_fullpath_files.py
#
# Name: Laura Tung
#
# Usage: python select_random_subset_fullpath_files.py <full_set_datasets_file> <subset_size>
#
# <full_set_datasets_file>: file listing the datasets (full-path files) in the full set.
# <subset_size>: desired size of the subset (i.e. number of datasets desired).

import sys
import numpy as np
import pandas as pd

if __name__ == "__main__":

    datasets_file = sys.argv[1]
    num_datasets = int(sys.argv[2])

    df = pd.read_csv(datasets_file, header=None, sep=' ', engine='python')
    print(df.info())

    selected_df = df.sample(n = num_datasets, random_state=1)

    outfile = "selected_subset_datasets_fullpath"
    pd.DataFrame.to_csv(selected_df, path_or_buf=outfile, sep=' ', header=False, index=False)
    np.savetxt("selected_subset_index", list(selected_df.index), fmt='%d')

max_stars_repo_name: Kingsford-Group/hierrepsetselection

max_issues_repo_name: Kingsford-Group/hierrepsetselection

max_forks_repo_name: Kingsford-Group/hierrepsetselection


text:
'''
This file includes util functions required while building neural network.
'''
import sys
import numpy as np
np.random.seed(1024)

def convert_to_1D(tensor):
    '''
    Converts a numpy array into an array of 1 dimension.
    :param tensor: input numpy array
    :return: 1D array
    '''
    return np.ravel(tensor)

def checkXandY(x,y):
    '''
    checks the number of points in two numpy array.
    :param x: numpy array
    :param y: numpy array
    :return: None
    '''
    assert (x.shape[0] == y.shape[0]), "Number of data points not equal"

def change_labels(ytrain):
    '''
    change the labels according to the data required by neural network.
    Example: ytrain = [1,0,1]
    change it to: [[0,1],[1,0],[0,1]]
    :param ytrain: numpy array of 1D dimension.
    :return: numpy array where the class label is triggered.
    '''

    ytrain = ytrain.ravel()
    numclass = len(np.unique(ytrain))
    changed_ytrain = []
    for label in ytrain:
        temp = np.zeros(numclass)
        temp[int(label)] =1
        changed_ytrain.append(temp)

    return np.array(changed_ytrain)

def generate_batches(data_size,batch_size):
    '''
    Generate batches of data based on the batch size.
    :param data_size: total number of data points.
    :param batch_size: batch size of type int
    :return:
    '''
    assert (isinstance(batch_size, int)), 'batch_size should be of type int, {0} not supported.'.format(type(batch_size))
    assert (batch_size > 0), 'batch_size should be greater than zero.'
    ind = np.array([i for i in range(data_size)])
    np.random.shuffle(ind)
    temp = []
    for i in range(0, data_size, batch_size):
        temp.append( ind[i:i + batch_size])
    return np.array(temp)

def progress_bar(curr_ite, total, total_epoch, curr_epoch,prefix = '', decimals = 0, barLength = 50):
    '''
    Prints the progress of a for loop.
    I have used the modified version of Greenstick code on stackoverflow.
    Link to the code: http://stackoverflow.com/questions/3173320/text-progress-bar-in-the-console?answertab=votes#tab-top
    :param curr_ite: current iteration
    :param total: total iterations
    :param total_epoch: total epochs to be conducted
    :param curr_epoch: current epoch
    :param prefix: not required(default Progress)
    :param decimals: positive number of decimals in percent complete (Int)
    :param barLength: length of the bar
    :return: None
    '''
    formatStr = "{0:." + str(decimals) + "f}"
    percent = formatStr.format(100 * (curr_ite / float(total)))
    filledLength = int(round(barLength * curr_ite / float(total)))
    bar = '>' * filledLength + '-' * (barLength - filledLength)
    sys.stdout.write('\repoch %s/%s %s %s %s%s' % (curr_epoch,total_epoch,prefix, bar, percent, '%')),
    if curr_ite == total:
        sys.stdout.write('\n')
    sys.stdout.flush()


text:
#!/usr/bin/env python3.6
"""
Multi gpu version, 2018.02.26 23:00 copied from mps_gpu_lowmem.py
02.28:  Single-GPU chunk-version is finished, and the file is renamed to mps_chunks.py
03.05:  Multiple-GPU version is finished.
03.06:  nll_history[] is added.
03.06:  some variables have been renamed, MPS.init() has been re-organized in a more elegant way,           which is also easy for saving/loading
03.06:  saving/loading using pickle.
        saving is carried out in training()
        loading fills the mps object, notice that after loading, self.broadcasting needs to be called to initialize data on each gpu.
03.10:  ztmp, _gpu_queue are redundant and removed
        altered idx_2site to ByteTensor, used as mask
        added self._tensortype to unify FloatTensor and DoubleTensor
03.10:  Removed unnecessary clone of self.newcumu
        Assigned self.newcumu and self.gradients to None after each usage
03.10:  Added sampling method
03.18:  Added attribute grad_mode
04.18:  Added attribute n_step
"""
#
#{{{ import
from __future__ import print_function
import numpy as np
from numpy.random import rand
import torch
import time
import sys,os
from multiprocessing.dummy import Pool as ThreadPool
import queue
import pickle
#}}}
#{{{ class MPS:
class MPS:
    """
    Matrix product states for generative modeling.
    """
    def __init__(self,data,batch_size,gpu_ids=[],use_double=False,init_bondim=2,data_test=[],fname="SGDresult",chunk_size=2000):
        """
        """
# Principles for storage: 1. variables for reading on GPU and variables for writing on CPU 2. Main thread always controls CPU and only one GPU for doing SVD and merging, all GPUs are used only in spawned threads.
# On CPU: self.tensors, self.cumulants, self.psis, self.psis_test, self.bondims, self.merged_tensor
# On GPU: self.data, self.data_test, self.merged_tensor_BC
# number of GPU cannot be changed
# batchsize & chunksize can be changed
        self._gpu_ids = gpu_ids # a list of gpu ids to use
        self._ngpu = len(gpu_ids)
        self.main_gpuid = self._gpu_ids[0] # id of gpu used for the main thread
        if use_double:
            self._tensortype = torch.DoubleTensor
            self.epsi = 1e-300
        else:
            self._tensortype = torch.FloatTensor
            self.epsi = 1e-26

        self.n=data.shape[1] # number of spins
        self.m=data.shape[0] # number of training samples
        self.m_test = len(data_test)
        assert self.m%len(gpu_ids)==0
        assert self.m_test%len(gpu_ids)==0
        self.m_gpu = self.m // len(gpu_ids)
        if self.m_test > 0:
            self.m_test_gpu = self.m_test // len(gpu_ids)

        self.data_cpu = data.type(torch.LongTensor).cpu()
        if(self.m_test > 0):
            self.data_test_cpu = data_test.type(torch.LongTensor).cpu()

        self.training_method=1 # 1: cumulants, default 2: slower one
        self.cutoff = 0.
        self.learning_rate = 0.01
        self.n_step = 1
        self.verbose = True
        self.converge_crit=0.01
        self.loop=0

        self.batch_size = self.m
        self.set_batchsize(batch_size)
        self.set_chunksize(chunk_size)
        self._build_idx()
        self.distribute()

        self.min_bondim = 2
        self.max_bondim = 200
        bondims = np.asarray([2**min(i,self.n-1-i,np.log2(init_bondim)) for i in range(self.n)],np.int16)
        self.bondims = torch.from_numpy(bondims)
        self.bondims.clamp_(self.min_bondim, self.max_bondim)
        self.bondims[-1] = 1

        self.randgn = np.random.RandomState(1996)
        self.seqSGD = False
        # Random initialization
        self.tensors= [torch.randn([self.bondims[i-1],2,self.bondims[i]]).type(self._tensortype) for i in range(self.n)]
        # # Prescribed 0^N
        # self.tensors = []
        # for i in range(self.n):
        #     tens = torch.zeros([self.bondims[i-1],2,self.bondims[i]]).type(self._tensortype)
        #     tens[0,0,0] = 1
        #     self.tensors.append(tens)
        # # Prescribed (0^N+1^N)
        # self.tensors = []
        # for i in range(self.n):
        #     tens = torch.zeros([self.bondims[i-1],2,self.bondims[i]]).type(self._tensortype)
        #     tens[0,0,0] = 1
        #     if i == 0 :
        #         tens[0,1,1] = 1
        #     elif i == self.n-1:
        #         tens[1,1,0] = 1
        #     else:
        #         tens[1,1,1] = 1
        #     self.tensors.append(tens)
        # # Prescribed (0+1)^N
        # self.tensors = []
        # for i in range(self.n):
        #     tens = torch.zeros([self.bondims[i-1],2,self.bondims[i]]).type(self._tensortype)
        #     tens[0,:,0] = 1
        #     self.tensors.append(tens)
        self.merged_tensor=self._tensortype() #self.merged_tensor will be on cpu and gid (in merge_bond() and rebuild_bond() for SVD), self.merged_tensor_BC will be the broadcasted version on each GPU
        self.tensor_bc=[]
        self.merged_tensor_BC=[] # broadcast version of merged_tensor, will be put on each GPU.

        self.grad_mode='plain'
        """
        plain: original GD;
        norm0.25: normalize with g.norm()/(g.numel()**0.25);
        norm0.5: normalize with g.norm()/(g.numel()**0.5);
        elementwise: elementwise divide by (g**2+5e-6)**0.5
        """
        self.gradients=None # storing gradients in each gpu, will be used in gradient_descent()
        self.current_bond=-1 # indicating that it's random
        self.batch_head=0
        self.going_right=True
        print("n=",self.n,"m=",self.m,"batch_size=",self.batch_size)

        sys.stdout.write("left_canonical ");sys.stdout.flush();t1=time.time()
        self.left_canonical()
        sys.stdout.write("%.2f s.\n"%(time.time()-t1))
        sys.stdout.write("Initiate NLL ");sys.stdout.flush();t1=time.time()
        self.recompute_psi()
        self.nll = self.compute_nll()
        self.nll_test = 0
        if self.m_test>0:
            self.recompute_psi_test()
            self.nll_test = self.compute_nll_test()
        print("#%d nll=%.3f\tnll_test=%.3f\tmax_bondim=%d\t%.2f Sec."%(self.loop,self.nll,self.nll_test,max(self.bondims),time.time()-t1))
        self.fname = fname+'.dat'
        self.nll_history=np.array([(self.loop,self.nll,self.nll_test)],dtype=[('loop','i'),('train','d'),('test','d')])
        with open(self.fname,'a') as fp:
            print('# nll_trn nll_tst max(bdm)\tduration\n%d %.3f %.3f %d'%(self.loop,self.nll,self.nll_test,max(self.bondims)),file=fp)
        with open('hyper.txt','a') as fp:
            print('max_bdm\tcut\tlr\tnstp\tbhead\tmbatch\tngpu\tgmode',file=fp)
#}}}
#{{{def set_batchsize(self,batch_size):
    def set_batchsize(self,batch_size):
        bsiz_old = self.batch_size
        self.batch_size_gpu = int(min(batch_size,self.m)//self._ngpu)
        self.batch_size = self.batch_size_gpu*self._ngpu
        self.learning_rate *= float(self.batch_size)/bsiz_old
        if hasattr(self, "chunk_size"):
            self.chunk_size = min(self.chunk_size, self.batch_size_gpu)
        # if hasattr(self, 'batch_idx') and self.batch_idx.numel() != self.batch_size_gpu + self.m_gpu:
        #     self.batch_idx = torch.from_numpy(np.concatenate((np.arange(self.m_gpu),np.arange(self.batch_size_gpu)))).type(torch.LongTensor)
#}}}
#{{{def shrink_chunk(self, chunk_size):
    def set_chunksize(self, chunk_size):
        self.chunk_size = min(chunk_size, self.batch_size_gpu)
        # self.chunk_nd = np.append(np.arange(0,self.batch_size_gpu,self.chunk_size),self.batch_size_gpu) #nodes for chunk
        # if self.m_test > 0:
        #     self.chunk_nd_test =np.append(np.arange(0,self.m_test_gpu,self.chunk_size),self.m_test_gpu)
#}}}
    def alterGPU(self, gpu_ids):
        self._gpu_ids = gpu_ids
        self._ngpu = len(gpu_ids)
        self.main_gpuid = gpu_ids[0]
        self.m_gpu = int(self.m // self._ngpu)
        if self.m_test > 0:
            self.m_test_gpu = int(self.m_test // self._ngpu)
        self.set_batchsize(self.batch_size)
#{{{def _build_idx(self):
    def _build_idx(self):
        print("Building the 2-site index:",end='',flush=True);tt=time.time()
        # self.idx_cpu=torch.arange(0,self.m).view(self._ngpu,-1).type(torch.LongTensor).pin_memory()
        # if self.m_test>0:
            # self.idx_test_cpu=torch.arange(0,self.m_test).view(self._ngpu,-1).type(torch.LongTensor).pin_memory()
        # self.idx_2sites_cpu[bond, i, j, :] is a vector of shape (self.m,) which indicates whether the sample meets sample[bond]==i and sample[bond+1]==j
        self.idx_2sites_cpu=torch.ByteTensor(self.n-1,2,2,self.m).pin_memory()
        for bond in range(self.n-1):
            for i in range(2):
                for j in range(2):
                    self.idx_2sites_cpu[bond,i,j,:] = ((self.data_cpu[:,bond]==i)*(self.data_cpu[:,bond+1]==j)).cpu()
        print("%.2f Sec."%(time.time()-tt))
#}}}
#{{{def _broadcast_idx(self):
    # def _broadcast_chunk_pt(self):
    #     assert self.batch_head + self.batch_size_gpu <= self.m_gpu
        # self._chunk_pt = torch.cuda.comm.broadcast(torch.LongTensor([0,self.chunk_size]), self._gpu_ids)
#}}}
#{{{def distribute(self):
    def distribute(self):
        self.data =torch.cuda.comm.scatter(self.data_cpu,self._gpu_ids)# declare data as long tensor for indexing purpuse, then broadcast it to all gpus
        # self.idx_gpu=torch.comm.broadcast(torch.arange(self.m_gpu+self.batch_size_gpu).type(torch.LongTensor), self._gpu_ids)
        if not hasattr(self,"idx_2sites_cpu"):
            self._build_idx()
        self.idx_2sites=torch.cuda.comm.scatter(self.idx_2sites_cpu,self._gpu_ids,dim=3)
        if self.m_test>0:
            self.data_test=torch.cuda.comm.scatter(self.data_test_cpu,self._gpu_ids)# declare data as long tensor for indexing purpuse, then broadcast it to all gpus
            # self.idx_test_gpu=torch.comm.broadcast(torch.arange(self.m_test_gpu).type(torch.LongTensor), self._gpu_ids)
        for gidx in range(self._ngpu):
            print("On GPU #%d, training data shape"%self._gpu_ids[gidx], tuple(self.data[gidx].shape), end=', ')
            if self.m_test>0:
                print("test data shape",tuple(self.data_test[gidx].shape))
            else:
                print("no test data")
        self.psis=torch.ones(self._ngpu,self.m_gpu).type(self._tensortype).pin_memory()
        #notice that I do not want to put psis into GPU, as it is convenient to share the memory across threads, and it is easy to compute NLL based on self.psis
        if(self.m_test >0):
            self.psis_test=torch.ones(self._ngpu,self.m_test_gpu).type(self._tensortype).pin_memory()
#}}}

#{{{def left_canonical(self):
    def left_canonical(self):
        if self.current_bond < 0:
            self.current_bond = 0
        # elif self.going_right:
        #     self.current_bond += 1
        self.going_right=True
        while self.current_bond < self.n-1:
            self.merge_bond()
            self.rebuild_bond(True)
            self.current_bond += 1
#}}}
#{{{def merge_bond(self):
    def merge_bond(self):
        #This function uses single gpu, self.main_gpuid, for doing contractions.
        assert(self.current_bond>=0 and self.current_bond<self.n)
        if self.current_bond == self.n-1:
            self.current_bond -= 1
        d=self.bondims[self.current_bond]
        self.merged_tensor=torch.mm(self.tensors[self.current_bond].cuda(self.main_gpuid).view(-1,d),\
                            self.tensors[self.current_bond+1].cuda(self.main_gpuid).view(d,-1))\
                            .view(self.bondims[self.current_bond-1],2,2,self.bondims[self.current_bond+1])
#}}}
#{{{def rebuild_bond(self,fix_bondim=False):
    def rebuild_bond(self,fix_bondim=False):
        """rebuild merged_tensor into two self.tensors{bl} and
        self.merged_tensor will be set to [] at the end of this funciton"""
        dl=self.bondims[self.current_bond-1]
        d=self.bondims[self.current_bond]
        dr=self.bondims[self.current_bond+1]
        self.merged_tensor=self.merged_tensor.view(dl*2,dr*2)
        try:
            [U,s,V]=torch.svd(self.merged_tensor)
        except:
            print("GPU SVD does not converge, sending it to CPU")
            [U,s,V]=torch.svd(self.merged_tensor.cpu())
            U=U.cuda(self.main_gpuid)
            s=s.cuda(self.main_gpuid)
            V=V.cuda(self.main_gpuid)
            print("CPU svd() finished")
        if not fix_bondim:
            d=torch.sum(s>self.cutoff*s[0])
            d=max(min(d,self.max_bondim),self.min_bondim)
        U=U[:,:d]
        S=torch.diag(s[:d])
        V=V[:,:d].t()
        if(self.going_right):
            V=torch.mm(S,V)
            V/=V.norm()
        else:
            U=torch.mm(U,S)
            U/=U.norm()
        self.tensors[self.current_bond]=U.contiguous().view(dl,2,d).cpu()
        self.bondims[self.current_bond]=d
        self.tensors[self.current_bond+1]=V.contiguous().view(d,2,dr).cpu()
        self.merged_tensor=None
        return s[:d]
#}}}

#{{{def train(self,n_loops):
    def train(self):
        assert self.grad_mode in ["plain","norm0.25","norm0.5","elementwise"]
        self.loop += 1
        t1=time.time()
        svd_time=0
        other_time=0
        if self.seqSGD:
            if self.batch_head > self.m_gpu-self.batch_size_gpu:
                self.batch_head = self.m_gpu-self.batch_size_gpu
        else:
            self.batch_head = self.randgn.randint(self.m_gpu - self.batch_size_gpu + 1)
        print("Batch head=%d, init cumulants..."%(self.batch_head), end='',flush=True)
        self.init_cumulants()
        print("%.2f Sec."%(time.time()-t1))
        self.going_right=False
        for bond in np.concatenate((np.arange(self.n-2,0,-1),np.arange(0,self.n-2))).astype(np.uint16):
            print("\r bond #%d / %d "%(bond,self.n), end='', flush=True)
            t0=time.time()
            self.current_bond=bond
            if(bond==0):
                self.going_right=True
            self.merge_bond()
            for s in range(self.n_step):
                self.gradient_descent()
                if s <self.n_step-1:
                    self.merged_tensor /= self.merged_tensor.norm()
            # Since there's only one step, no need to normalize merged_tensor
            other_time += time.time()-t0; t0=time.time()
            self.rebuild_bond(False)
            svd_time += time.time()-t0; t0=time.time()
            self.update_cumulants()
            other_time += time.time()-t0
        t0=time.time()
        sys.stdout.write("\r                                                       \r")
        self.recompute_psi()
        nll_new=self.compute_nll()
        if(self.m_test >0):
            self.recompute_psi_test()
            self.nll_test=self.compute_nll_test()
        other_time += time.time()-t0
        sys.stdout.write("\r                                                       \r")
        print("#%d batch_head=%d\tnll=%.3f\tnll_test=%.3f\tmax_bondim=%d\t%.2f Sec., SVD %.2f, other %.2f"%(self.loop, self.batch_head, nll_new,self.nll_test,max(self.bondims),time.time()-t1,svd_time,other_time))
        if abs(nll_new-self.nll) < self.converge_crit:
            print("Converged")
            not_converged=False
        else:
            not_converged=True
        self.nll=nll_new
        self._log(time.time()-t1)
        if self.seqSGD:
            self.batch_head += self.batch_size_gpu
            if self.batch_head >= self.m_gpu:
                self.batch_head = 0
        return not_converged
#}}}
#{{{def jobs_threading(self,func):
    def _jobs_threading(self,func):
        self._job_queue=queue.Queue()
        pool=ThreadPool(self._ngpu)
        jobs=[]
        for gidx in range(self._ngpu):
            jobs.append(gidx)
            self._job_queue.put(gidx)
        pool.map(func,jobs)
        pool.close()
        while(not self._job_queue.empty()):
            time.sleep(0.01)
        else:
            assert self._job_queue.empty()
#}}}
#{{{def _worker_init_cumulants(self,gidx):
    def _worker_init_cumulants(self,gidx):
        gpuid=self._gpu_ids[gidx]
        states = self.data[gidx][self.batch_head:self.batch_head+self.batch_size_gpu, self.current_bond]
        chunk_begin, chunk_end = (0, self.chunk_size)
        while chunk_begin < self.batch_size_gpu:
            state_loc=states[chunk_begin:chunk_end]
            self.cumulants[self.current_bond+1][gidx,chunk_begin:chunk_end,:,:]=\
                    (self.cumulants[self.current_bond][gidx,chunk_begin:chunk_end,:,:].cuda(gpuid) \
                    @ self.tensor_bc[gidx][:,state_loc,:].permute(1,0,2)).cpu()
            self.cumulants[self.current_bond+1][gidx,chunk_begin:chunk_end,:,:] /= \
                max(self.cumulants[self.current_bond+1][gidx,chunk_begin:chunk_end].abs().max(),1e-6)
            ### normalize the psi values, for saving the precison.
            ### It works together with adding small values for resolving the non psi issue
            chunk_begin = chunk_end
            chunk_end = min(self.batch_size_gpu, chunk_end+self.chunk_size)
        self._job_queue.get()
#}}}
#{{{ def init_cumulants(self,batch):
    def init_cumulants(self):
        """
        self.cumulants are created here for each batch, and stored in cpu
        size of self.cumulants is [cumulants[0],cumulants[1],...,cumulants[n]], each of which is a vector (stored as torch.tensor) of length batch_size.
        """
        self.cumulants=[] # its length is self.n
        self.cumulants.append(torch.ones(self._ngpu,self.batch_size_gpu,1,1).type(self._tensortype).pin_memory())
        for site in range(self.n-2): # n-2 for two-site update, notice that this would be n-1 for single-site update
            self.current_bond=site
            self.cumulants.append(torch.ones(self._ngpu,self.batch_size_gpu,1,self.bondims[site]).type(self._tensortype).pin_memory())
            self.tensor_bc=torch.cuda.comm.broadcast(self.tensors[site],self._gpu_ids)
            self._jobs_threading(self._worker_init_cumulants)
        self.cumulants.append(torch.ones(self._ngpu,self.batch_size_gpu,1,1).type(self._tensortype).pin_memory())
        self.tensor_bc = None
#}}}
#{{{def _worker_update_cumulants(self,gidx):
    def _worker_update_cumulants(self,gidx):
        gpuid= self._gpu_ids[gidx]
        btmp = self.current_bond+1 if self.going_right else self.current_bond
        if self.going_right:
            states_loc = self.data[gidx][self.batch_head:self.batch_head+self.batch_size_gpu,self.current_bond]
        else:
            states_loc = self.data[gidx][self.batch_head:self.batch_head+self.batch_size_gpu,self.current_bond+1]
        chunk_begin, chunk_end = (0, self.chunk_size)
        while chunk_begin < self.batch_size_gpu:
            if(self.going_right):
                newcumu = self.cumulants[self.current_bond][gidx,chunk_begin:chunk_end].cuda(gpuid)\
                         @ (self.tensor_bc[gidx][:,states_loc[chunk_begin:chunk_end]].permute(1,0,2))
            else:
                newcumu = self.tensor_bc[gidx][:,states_loc[chunk_begin:chunk_end]].permute(1,0,2)\
                         @ self.cumulants[self.current_bond+1][gidx,chunk_begin:chunk_end].cuda(gpuid)
            newcumu /= max(newcumu.abs().max(),1e-6) ## normalize among this chunk
            self.cumulants[btmp][gidx,chunk_begin:chunk_end]=newcumu.cpu()
            chunk_begin = chunk_end
            chunk_end = min(self.batch_size_gpu, chunk_end+self.chunk_size)
        self._job_queue.get()
#}}}
#{{{def update_cumulants(self):
    def update_cumulants(self):
        """
        Update cumulants for after one tensor is learnt in the current batch.
        """
        if(self.going_right):
            self.tensor_bc=torch.cuda.comm.broadcast(self.tensors[self.current_bond],self._gpu_ids)
            self.cumulants[self.current_bond+1]=self._tensortype(self._ngpu,self.batch_size_gpu,1,self.bondims[self.current_bond]).pin_memory()
        else:
            self.tensor_bc=torch.cuda.comm.broadcast(self.tensors[self.current_bond+1],self._gpu_ids)
            self.cumulants[self.current_bond]=self._tensortype(self._ngpu,self.batch_size_gpu,self.bondims[self.current_bond],1).pin_memory()
        self._jobs_threading(self._worker_update_cumulants)
        self.tensor_bc = None
        # btmp = self.current_bond+1 if self.going_right else self.current_bond
        # self.cumulants[btmp]=self.newcumu_batch
        # self.newcumu_batch = None
        ### normalize the psi values, for saving the precison.
        ### It works together with adding small values for resolving the non psi issue
#}}}
#{{{def _worker_gradient_descent(self,gidx):
    def _worker_gradient_descent(self,gidx):
        gpuid=self._gpu_ids[gidx]
        chunk_begin, chunk_end = (0,self.chunk_size)
        local_idx_2sites_batch = self.idx_2sites[gidx][self.current_bond,:,:,self.batch_head:self.batch_head+self.batch_size_gpu]
        while chunk_begin < self.batch_size_gpu:
            cumu1=self.cumulants[self.current_bond][gidx,chunk_begin:chunk_end,:,:].cuda(gpuid)
            cumu2=self.cumulants[self.current_bond+1][gidx,chunk_begin:chunk_end,:,:].cuda(gpuid)
            converter  = torch.arange(chunk_end-chunk_begin).type(torch.LongTensor).cuda(gpuid)
            for i,j in [(i,j) for i in range(2) for j in range(2)]:
                idx = converter[local_idx_2sites_batch[i,j,chunk_begin:chunk_end]]
                if(idx.numel()==0):
                    continue
                lvecs=cumu1[idx]
                rvecs=cumu2[idx]
                psis = lvecs @ self.merged_tensor_BC[gidx][:,i,j,:].repeat(idx.numel(),1,1) @ rvecs
                psis += self.epsi  ### this is to resolve the zero psis issue, for the float-storage
                lvecs /= psis
                torch.addbmm(self.gradients[gidx][:,i,j,:],lvecs.permute(0,2,1),rvecs.permute(0,2,1),out=self.gradients[gidx][:,i,j,:])
            chunk_begin = chunk_end
            chunk_end = min(self.batch_size_gpu, chunk_end+self.chunk_size)
        self._job_queue.get()
#}}}
#{{{def gradient_descent(self,batch):
    def gradient_descent(self):
        self.merged_tensor_BC=torch.cuda.comm.broadcast(self.merged_tensor,self._gpu_ids)
        self.gradients = [torch.zeros(self.bondims[self.current_bond-1],2,2,self.bondims[self.current_bond+1]).type(self._tensortype).cuda(gid) for gid in self._gpu_ids]
        self._jobs_threading(self._worker_gradient_descent)
        gradients = torch.cuda.comm.reduce_add(self.gradients, self.main_gpuid).cuda(self.main_gpuid)/self.batch_size-self.merged_tensor
        if self.grad_mode=='plain':
            pass
        elif self.grad_mode[:4]=='norm':
            if self.grad_mode=='norm0.25':
                gnorm = gradients.norm()/(gradients.numel()**0.25) # saddle killer?
            elif self.grad_mode=='norm0.5':
                gnorm = gradients.norm()/(gradients.numel()**0.5)
            if gnorm < 1:
                gradients /= gnorm
        elif self.grad_mode=='elementwise':
            gradients /= (gradients**2+5e-6)**0.5
        self.merged_tensor += gradients*2.0*self.learning_rate
        self.gradients = None
        self.merged_tensor_BC = None
#}}}

#{{{def compute_nll(self):
    def compute_nll(self):
        return -torch.log(self.psis.abs()).mean()*2.0
#}}}
#{{{    def compute_nll_test(self):
    def compute_nll_test(self):
        return -torch.log(self.psis_test.abs()).mean()*2.0
#}}}

#{{{def _worker_recompute_psi(self,gidx):
    def _worker_recompute_psi(self,gidx):
        gpuid=self._gpu_ids[gidx]
        chunk_begin,chunk_end = (0,self.chunk_size)
        while chunk_begin < self.m_gpu:
            state=self.data[gidx][chunk_begin:chunk_end,:]
            psis=torch.ones(chunk_end-chunk_begin,1,1).type(self._tensortype).cuda(gpuid)
            for i in range(self.n):
                psis @= self.tensors[i].cuda(gpuid)[:,state[:,i],:].permute(1,0,2)
            self.psis[gidx,chunk_begin:chunk_end]=psis.view(-1,)
            chunk_begin = chunk_end
            chunk_end = min(chunk_end+self.chunk_size, self.m_gpu)
        self._job_queue.get()
#}}}
#{{{def recompute_psi(self):
    def recompute_psi(self):
        self._jobs_threading(self._worker_recompute_psi)
#}}}
#{{{def _worker_recompute_psi_test(self,gidx):
    def _worker_recompute_psi_test(self,gidx):
        gpuid=self._gpu_ids[gidx]
        chunk_begin,chunk_end = (0,min(self.chunk_size,self.m_test_gpu))
        while chunk_begin < self.m_test_gpu:
            state=self.data_test[gidx][chunk_begin:chunk_end,:]
            psis=torch.ones(chunk_end-chunk_begin,1,1).type(self._tensortype).cuda(gpuid)
            for i in range(self.n):
                psis @= self.tensors[i].cuda(gpuid)[:,state[:,i],:].permute(1,0,2)
            self.psis_test[gidx,chunk_begin:chunk_end]=psis.view(-1,)
            chunk_begin = chunk_end
            chunk_end = min(chunk_end+self.chunk_size, self.m_test_gpu)
        self._job_queue.get()
#}}}
#{{{def recompute_psi_test(self,batch):
    def recompute_psi_test(self):
        assert self.m_test >0
        self._jobs_threading(self._worker_recompute_psi_test)
#}}}
#{{{def _log(self):
    def _log(self, duration=0):
        self.nll_history = np.append(self.nll_history,np.array((self.loop,self.nll,self.nll_test),dtype=self.nll_history.dtype))
        with open(self.fname,"a") as fout:
            print("%d %.3f %.3f %d\t\t%.1f"%(self.loop,self.nll,self.nll_test,max(self.bondims),duration),file=fout,flush=True)
        with open('hyper.txt','a') as fp:
            #max_bdm\tcut\tlr\tnstp\tbhead\tmbatch\tngpu\tgmode
            print("%d\t%.2g\t%.3g\t%d\t%d\t%d\t%d\t%s"%\
            (self.max_bondim,self.cutoff,self.learning_rate,self.n_step,self.batch_head,self.batch_size,self._ngpu,self.grad_mode),file=fp)
#}}}
#{{{def generate_sample(self, stat=None, givn_mask=None):
    def generate_sample(self, stat=None, givn_mask=None, gid=0):
        """
        If stat is None, generate from scratch, return a direct sample;
        Else stat will be cloned and masked by givn_mask, sites where givn_mask==True will be clamped before sampling, return a conditional sample
        """
        gpu = self._gpu_ids[gid]
        # print("On GPU #%d"%gpu)
        if stat is None or givn_mask is None or givn_mask.any()==False:
            assert self.current_bond == self.n - 1
            state = np.empty((self.n,), dtype=np.uint8)
            vec = torch.ones(1).type(self._tensortype).cuda(gpu)
            for p in np.arange(self.n)[::-1]:
                vec_act = self.tensors[p][:, 1].cuda(gpu) @ vec
                nom = vec_act.norm()
                if rand() < nom**2:
                    state[p] = 1
                    vec = vec_act/nom
                else:
                    state[p] = 0
                    vec = self.tensors[p][:, 0].cuda(gpu) @ vec
                    vec /= vec.norm()
            return state
        if isinstance(stat, np.ndarray):
            state = stat.astype(np.uint8)
        else:
            state = stat.cpu().numpy().astype(np.int8)
        state[givn_mask==False] = -1
        # givn_mask = givn_mask.copy()

        given_idx = givn_mask.nonzero()[0]
        p_given_rightmost = given_idx.max()
        p_given_leftmost  = given_idx.min()

        p_unnorm = self.current_bond
        if self.current_bond != self.n-1 and self.going_right:
                p_unnorm += 1
        if givn_mask[p_unnorm] == False:
            targ_p_unnorm = given_idx[np.argmin(abs(given_idx-p_unnorm))]
            if p_unnorm > targ_p_unnorm:
                self.going_right = False
                while p_unnorm > targ_p_unnorm:
                    self.current_bond = p_unnorm-1
                    self.merge_bond()
                    self.rebuild_bond(True)
                    p_unnorm -= 1
                    print('Canonicalization: unnormalized site = #%d'%p_unnorm, end='\r', flush=True)
            elif p_unnorm < targ_p_unnorm:
                self.going_right = True
                while p_unnorm < targ_p_unnorm:
                    self.current_bond = p_unnorm
                    self.merge_bond()
                    self.rebuild_bond(True)
                    p_unnorm += 1
                    print('Canonicalization: unnormalized site = #%d'%p_unnorm, end='\r', flush=True)
        else:
            # print(givn_mask[bd],givn_mask[bd+1])
            print('No extra canonicalization!',end='')
        print('')
        # From now on matrices are fixed

        # print("Locating p_intermid_ungiven_leftmost")
        p = p_given_leftmost
        while p<self.n and givn_mask[p]:
            p += 1
        p_intermid_ungiven_leftmost = p
        # print('Given leftmost',p_given_leftmost, ', p_intermid_ungiven_leftmost', p_intermid_ungiven_leftmost, ', rightmost', p_given_rightmost)

        mid_mat_fromleft = self.tensors[p_given_leftmost][:,state[p_given_leftmost]].cuda(gpu)
        mid_mat_fromleft /= mid_mat_fromleft.norm()
        for p in range(p_given_leftmost+1, p_intermid_ungiven_leftmost):
            mid_mat_fromleft @= self.tensors[p][:,state[p]].cuda(gpu)
            mid_mat_fromleft /= mid_mat_fromleft.norm()

        if p_intermid_ungiven_leftmost < p_given_rightmost+1:
            right_vecs = np.empty(self.n, dtype=object)
            # right_vecs[p] is the part of site>=p and its right that contributes to the marginal of some sites <p
            p = p_given_rightmost
            tens = self.tensors[p][:,state[p]].cuda(gpu)
            tens /= tens.norm()
            right_vecs[p] = (tens@tens.t())
            # right_vecs[p] /= np.trace(right_vecs[p])
            p -= 1
            while p > p_intermid_ungiven_leftmost:
                if givn_mask[p]:
                    tens = self.tensors[p][:,state[p]].cuda(gpu)
                    right_vecs[p] = tens @ right_vecs[p+1] @ tens.t()
                else:
                    tens = self.tensors[p].cuda(gpu)
                    right_vecs[p] = tens.view(-1,tens.shape[-1]) @ right_vecs[p+1]
                    right_vecs[p] = right_vecs[p].view(tens.shape[0],-1) @ (tens.view(tens.shape[0],-1).t())
                # right_vecs[p+1] = right_vecs[p+1].cpu()
                right_vecs[p+1] = right_vecs[p+1].cuda(gpu)
                right_vecs[p] /= right_vecs[p].trace()
                p -= 1
            # print("\nSampling intermediate bits")
            p = p_intermid_ungiven_leftmost
            while p < p_given_rightmost:
                # print('\r %d/%d'%(p,self.n), end="", flush=True)
                tens = self.tensors[p].cuda(gpu)
                if givn_mask[p]==False:
                    # print("Current bit:",p,"given:F","left_vec.shape", left_vec.shape, "matrix.shape", self.tensors[p].shape)
                    right_vec = right_vecs[p+1].cuda(gpu)
                    # prob_marg = tens.view(-1,tens.shape[-1]) @ right_vec
                    # prob_marg = prob_marg.view(tens.shape[0],-1) @ (tens.view(tens.shape[0],-1).t())
                    # prob_marg = (mid_mat_fromleft @ prob_marg @ mid_mat_fromleft.t()).trace()
                    # mid_mat_act = mid_mat_fromleft @ tens[:,1].cuda(gpu)
                    # prob_actv = (mid_mat_act @ right_vec @ mid_mat_act.t()).trace()
                    # if rand()<prob_actv/prob_marg:
                    #     state[p] = 1
                    #     mid_mat_fromleft = mid_mat_act
                    # else:
                    #     state[p] = 0
                    #     mid_mat_fromleft @= tens[:,0].cuda(gpu)
                    mid_mat_1 = mid_mat_fromleft @ tens[:,1].cuda(gpu)
                    mid_mat_0 = mid_mat_fromleft @ tens[:,0].cuda(gpu)
                    prob_1 = (mid_mat_1 @ right_vec @ (mid_mat_1.t())).trace()
                    prob_0 = (mid_mat_0 @ right_vec @ (mid_mat_0.t())).trace()
                    if rand() < prob_1/(prob_0+prob_1):
                        state[p] = 1
                        mid_mat_fromleft = mid_mat_1
                    else:
                        state[p] = 0
                        mid_mat_fromleft = mid_mat_0
                    # givn_mask[p] = True
                else:
                    # print("Current bit:",p,"given:T","left_vec.shape", left_vec.shape, "matrix.shape", self.tensors[p].shape)
                    mid_mat_fromleft @= tens[:,state[p]].cuda(gpu)
                mid_mat_fromleft /= mid_mat_fromleft.norm()
                right_vecs[p+1]=None
                p += 1
            mid_mat_fromleft @= self.tensors[p][:,state[p]].cuda(gpu)
            mid_mat_fromleft /= mid_mat_fromleft.norm()
            p += 1
            # assert givn_mask[p_given_leftmost:p_given_rightmost+1].all()
            # print('')
        for p in range(p_given_rightmost+1, self.n):
            vec_act = mid_mat_fromleft @ self.tensors[p][:,1].cuda(gpu)
            nom = vec_act.norm()
            if rand() < nom**2: #activate
                state[p] = 1
                mid_mat_fromleft = vec_act/nom
            else: #keep 0
                state[p] = 0
                mid_mat_fromleft @= self.tensors[p][:,0].cuda(gpu)
                mid_mat_fromleft /= mid_mat_fromleft.norm()
        for p in np.arange(p_given_leftmost)[::-1]:
            vec_act = self.tensors[p][:,1].cuda(gpu) @ mid_mat_fromleft
            nom = vec_act.norm()
            if rand() < nom**2:
                state[p] = 1
                mid_mat_fromleft = vec_act/nom
            else:
                state[p] = 0
                mid_mat_fromleft = self.tensors[p][:,0].cuda(gpu) @ mid_mat_fromleft
                mid_mat_fromleft /= mid_mat_fromleft.norm()
        assert (state!=-1).all()
        return state
#}}}
#{{{def __getstate__(self):
    def __getstate__(self):
        """ Retuen a dictionary of state values to be pickled """
        exclusion=["tensors","data_cpu","data_test_cpu","data","data_test","idx_gpu","cumulants","gradients","_job_queue","idx_2sites"]
        mydict={}
        for key in self.__dict__.keys():
            if key not in exclusion:
                mydict[key] = self.__dict__[key]
        return mydict
#}}}
#{{{def save(self, fsave_name)
    def save(self, fsave_name=None):
        orig_dir = os.getcwd()
        if fsave_name:
            if fsave_name[0]=='/':
                dit = fsave_name
            else:
                dit = orig_dir+'/'+fsave_name
        else:
            dit = orig_dir+'/L%d'%self.loop
        while os.path.exists(dit):
            dit = dit+'_'
        os.mkdir(dit)
        os.chdir(dit)
        with open('mps.pickle','wb') as fsave:
            tsave=time.time()
            sys.stdout.write("saving mps object to %s ..."%dit);sys.stdout.flush()
            pickle.dump(self,fsave)
            sys.stdout.write("%.2f Sec.\n"%(time.time()-tsave))
        np.savez_compressed('tensors.npz',*[t.numpy() for t in self.tensors])
        if self.m_test>0:
            np.savez_compressed('../data.npz', train=self.data_cpu.numpy().astype(np.uint8), test=self.data_test_cpu.numpy().astype(np.uint8))
        else:
            np.savez_compressed('../data.npz', train=self.data_cpu.numpy().astype(np.uint8), test=None)
        os.chdir(orig_dir)
#}}}

def mpsLoad(path):
    original = os.getcwd()
    os.chdir(path)
    print("Loading",path,end='...',flush=True)
    f = open('mps.pickle','rb')
    mps = pickle.load(f)
    dat = np.load("../data.npz")
    mps.data_cpu=torch.from_numpy(dat['train'].astype(np.int64))
    if dat['test'].size > 1:
        mps.data_test_cpu=torch.from_numpy(dat['test'].astype(np.int64))
    mps.tensors=[]
    tens = np.load('tensors.npz')
    for i in range(mps.n):
        mps.tensors.append(torch.from_numpy(tens['arr_%d'%i]))
    f.close()
    print("Done.")
    os.chdir(original)
    return mps

#{{{if __name__ == "__main__":
if __name__ == "__main__":
    mps = mpsLoad(sys.argv[1])

max_stars_repo_name: congzlwag/UnsupGenModbyMPS

max_issues_repo_name: congzlwag/UnsupGenModbyMPS

max_forks_repo_name: congzlwag/UnsupGenModbyMPS

reason: import numpy,from numpy