<commit_before>import os
import sys
import json
import bz2
import seaborn
import pandas as pd
from unicode_codes import EMOJI_UNICODE
from twitter_search_funcs import *
from timeit import default_timer as timer

data_path = "/your/data/path/archive-twitter-2016-08/"

# Character to match
match = EMOJI_UNICODE[':pistol:']

# Counters
counter_total_tweets = 0
counter_total_match = 0
counter_total_before = 0
counter_total_after = 0
counterdict_before = {}
counterdict_after = {}
counterdict_lang = {}

# Main search loop
start_t = timer()
for day in range(31):
    day_str = "{:02d}".format(day + 1)
    for hour in range(24):
        hour_str = "{:02d}".format(hour)
        file_path = os.path.join(data_path, day_str, hour_str)
        files = os.listdir(file_path)
        for i, f in enumerate(files):
            progress(i + hour * len(files),
                     len(files) * 24,
                     suffix="Searching Day {}, Hour {}".format(day_str, hour_str))
            fbz = bz2.BZ2File(os.path.join(file_path, f), 'rb')
            try:
                fdec = fbz.read()
            except IOError:
                continue
            finally:
                fbz.close()
            fdecutf = fdec.decode('utf-8')

            for t in fdecutf.split('\n'):
                try:
                    tweet = json.loads(t)
                except (ValueError):
                    continue
                if 'delete' in tweet.keys():
                    continue
                counter_total_tweets += 1
                if match in tweet['text']:
                    counter_total_match += 1
                    result = find_context(tweet['text'], match)

                    if result[0] in EMOJI_UNICODE.values():
                        counter_total_before += 1

                        if result[0] in counterdict_before.keys():
                            counterdict_before[result[0]] += 1
                        else:
                            counterdict_before[result[0]] = 1

                    if result[2] in EMOJI_UNICODE.values():
                        counter_total_after += 1

                        if result[2] in counterdict_after.keys():
                            counterdict_after[result[2]] += 1
                        else:
                            counterdict_after[result[2]] = 1

                    try:
                        if tweet['lang'] in counterdict_lang.keys():
                            counterdict_lang[tweet['lang']] += 1
                        else:
                            counterdict_lang[tweet['lang']] = 1
                    except KeyError:
                        continue
end_t = timer()

# Print outputs
print("Elapsed Time    : {:.2f} min".format((end_t - start_t) / 60))
print("Total Tweets    : {:d}".format(counter_total_tweets))
print("Total Matches   : {:d}".format(counter_total_match))
print("Total w/ Before : {:d}".format(counter_total_before))
print("Total w/ After  : {:d}".format(counter_total_after))

# Convert output to dataframe
df_before = pd.DataFrame(list(counterdict_before.items()), columns=['Emoji', 'CountBefore'])
df_after = pd.DataFrame(list(counterdict_after.items()), columns=['Emoji', 'CountAfter'])
df_lang = pd.DataFrame(list(counterdict_lang.items()), columns=['Lang', 'Count'])

# Merge before and after dataframes
df_all = pd.merge(df_before, df_after, on='Emoji', how='outer')

df_all.sort_values('CountBefore', ascending=False).head()
df_all.sort_values('CountAfter', ascending=False).head()
df_lang.sort_values('Count', ascending=False).head()
df_all.sort_values('CountBefore', ascending=False).head(20).plot.bar(y='CountBefore')
df_all.sort_values('CountAfter', ascending=False).head(20).plot.bar(y='CountAfter')

# Export results as CSV files
df_all.to_csv("./alldata.csv")
df_lang.to_csv("./langdata.csv")
<commit_msg>Remove unused imports, replace star import with explicit, and sort<commit_after>#!/usr/bin/env python
# encoding: utf-8
import bz2
import json
import os
import pandas as pd
from timeit import default_timer as timer
from twitter_search_funcs import find_context, progress
from unicode_codes import EMOJI_UNICODE

data_path = "/your/data/path/archive-twitter-2016-08/"

# Character to match
match = EMOJI_UNICODE[':pistol:']

# Counters
counter_total_tweets = 0
counter_total_match = 0
counter_total_before = 0
counter_total_after = 0
counterdict_before = {}
counterdict_after = {}
counterdict_lang = {}

# Main search loop
start_t = timer()
for day in range(31):
    day_str = "{:02d}".format(day + 1)
    for hour in range(24):
        hour_str = "{:02d}".format(hour)
        file_path = os.path.join(data_path, day_str, hour_str)
        files = os.listdir(file_path)
        for i, f in enumerate(files):
            progress(i + hour * len(files),
                     len(files) * 24,
                     suffix="Searching Day {}, Hour {}".format(day_str, hour_str))
            fbz = bz2.BZ2File(os.path.join(file_path, f), 'rb')
            try:
                fdec = fbz.read()
            except IOError:
                continue
            finally:
                fbz.close()
            fdecutf = fdec.decode('utf-8')

            for t in fdecutf.split('\n'):
                try:
                    tweet = json.loads(t)
                except (ValueError):
                    continue
                if 'delete' in tweet.keys():
                    continue
                counter_total_tweets += 1
                if match in tweet['text']:
                    counter_total_match += 1
                    result = find_context(tweet['text'], match)

                    if result[0] in EMOJI_UNICODE.values():
                        counter_total_before += 1

                        if result[0] in counterdict_before.keys():
                            counterdict_before[result[0]] += 1
                        else:
                            counterdict_before[result[0]] = 1

                    if result[2] in EMOJI_UNICODE.values():
                        counter_total_after += 1

                        if result[2] in counterdict_after.keys():
                            counterdict_after[result[2]] += 1
                        else:
                            counterdict_after[result[2]] = 1

                    try:
                        if tweet['lang'] in counterdict_lang.keys():
                            counterdict_lang[tweet['lang']] += 1
                        else:
                            counterdict_lang[tweet['lang']] = 1
                    except KeyError:
                        continue
end_t = timer()

# Print outputs
print("Elapsed Time    : {:.2f} min".format((end_t - start_t) / 60))
print("Total Tweets    : {:d}".format(counter_total_tweets))
print("Total Matches   : {:d}".format(counter_total_match))
print("Total w/ Before : {:d}".format(counter_total_before))
print("Total w/ After  : {:d}".format(counter_total_after))

# Convert output to dataframe
df_before = pd.DataFrame(list(counterdict_before.items()), columns=['Emoji', 'CountBefore'])
df_after = pd.DataFrame(list(counterdict_after.items()), columns=['Emoji', 'CountAfter'])
df_lang = pd.DataFrame(list(counterdict_lang.items()), columns=['Lang', 'Count'])

# Merge before and after dataframes
df_all = pd.merge(df_before, df_after, on='Emoji', how='outer')

df_all.sort_values('CountBefore', ascending=False).head()
df_all.sort_values('CountAfter', ascending=False).head()
df_lang.sort_values('Count', ascending=False).head()
df_all.sort_values('CountBefore', ascending=False).head(20).plot.bar(y='CountBefore')
df_all.sort_values('CountAfter', ascending=False).head(20).plot.bar(y='CountAfter')

# Export results as CSV files
df_all.to_csv("./alldata.csv")
df_lang.to_csv("./langdata.csv")
<|endoftext|>


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

"""Tests for tensor2tensor.envs.env_problem."""

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

import os
from gym.spaces import Box
from gym.spaces import Discrete
import numpy as np
from tensor2tensor.data_generators import generator_utils
from tensor2tensor.data_generators import problem
from tensor2tensor.envs import env_problem
from tensor2tensor.envs import env_problem_utils
from tensor2tensor.layers import modalities
import tensorflow as tf

class EnvProblemTest(tf.test.TestCase):

  def setUp(self):
    self.tmp_dir = os.path.join(tf.test.get_temp_dir(), "tmp_dir")
    tf.gfile.MakeDirs(self.tmp_dir)

  def tearDown(self):
    tf.gfile.DeleteRecursively(self.tmp_dir)

  def test_setup(self):
    ep = env_problem.EnvProblem(base_env_name="CartPole-v0", batch_size=5)
    # Checks that environments were created and they are `batch_size` in number.
    ep.assert_common_preconditions()

    # Expectations on the observation space.
    observation_space = ep.observation_space
    self.assertTrue(isinstance(observation_space, Box))
    self.assertEqual(observation_space.shape, (4,))
    self.assertEqual(observation_space.dtype, np.float32)

    # Expectations on the action space.
    action_space = ep.action_space
    self.assertTrue(isinstance(action_space, Discrete))
    self.assertEqual(action_space.shape, ())
    self.assertEqual(action_space.dtype, np.int64)
    self.assertEqual(ep.num_actions, 2)

    # Reward range is infinite here.
    self.assertFalse(ep.is_reward_range_finite)

  def test_reward_range(self):
    # Passing reward_range=None means take the reward range of the underlying
    # environment as the reward range.
    ep = env_problem.EnvProblem(
        base_env_name="FrozenLake-v0", batch_size=5, reward_range=None)
    ep.assert_common_preconditions()

    # Assert reward range is finite here.
    self.assertTrue(ep.is_reward_range_finite)

    # Assert that it is as expected of the underlying environment, since reward_
    self.assertEqual(0, ep.reward_range[0])
    self.assertEqual(1, ep.reward_range[1])

  def test_default_processed_rewards_discrete(self):
    # This differs in the above because it has a Tuple observation space.
    ep = env_problem.EnvProblem(
        base_env_name="KellyCoinflip-v0", batch_size=5, reward_range=None)
    ep.assert_common_preconditions()

    # Assert reward range is finite here.
    self.assertTrue(ep.is_reward_range_finite)

    # Assert that it is as expected of the underlying environment.
    reward_range = ep.reward_range
    self.assertEqual(0, reward_range[0])
    self.assertEqual(ep._envs[0].maxWealth, reward_range[1])

    # Check that the processed rewards are discrete.
    self.assertTrue(ep.is_processed_rewards_discrete)

    # Assert on the number of rewards.
    self.assertEqual(ep.num_rewards, reward_range[1] - reward_range[0] + 1)

  def test_interaction_with_env(self):
    batch_size = 5
    reward_range = (-1, 1)
    ep = env_problem.EnvProblem(
        base_env_name="KellyCoinflip-v0",
        batch_size=batch_size,
        reward_range=reward_range)

    # Resets all environments.
    ep.reset()

    # Let's play a few steps.
    nsteps = 100
    num_trajectories_completed = 0
    num_timesteps_completed = 0
    # If batch_done_at_step[i] = j then it means that i^th env last got done at
    # step = j.
    batch_done_at_step = np.full(batch_size, -1)
    for i in range(nsteps):
      # Sample batch_size actions from the action space and stack them (since
      # that is the expected type).
      actions = np.stack([ep.action_space.sample() for _ in range(batch_size)])

      _, _, dones, _ = ep.step(actions)

      # Do the book-keeping on number of trajectories completed and expect that
      # it matches ep's completed number.

      num_done = sum(dones)
      num_trajectories_completed += num_done

      self.assertEqual(num_trajectories_completed,
                       len(ep.trajectories.completed_trajectories))

      # Get the indices where we are done ...
      done_indices = env_problem_utils.done_indices(dones)

      # ... and reset those.
      ep.reset(indices=done_indices)

      # If nothing got done, go on to the next step.
      if done_indices.size == 0:
        # i.e. this is an empty array.
        continue

      # See when these indices were last done and calculate how many time-steps
      # each one took to get done.
      num_timesteps_completed += sum(i + 1 - batch_done_at_step[done_indices])
      batch_done_at_step[done_indices] = i

      # This should also match the number of time-steps completed given by ep.
      num_timesteps_completed_ep = sum(
          ct.num_time_steps for ct in ep.trajectories.completed_trajectories)
      self.assertEqual(num_timesteps_completed, num_timesteps_completed_ep)

    # Reset the trajectories.
    ep.trajectories.reset_batch_trajectories()
    self.assertEqual(0, len(ep.trajectories.completed_trajectories))

  def read_tfrecord_dataset(self, filenames, ep):
    # Read the dataset at `filenames` into a tf.data.Dataset and returns the
    # number of time-steps (just the number of records in the dataset) and the
    # number of trajectories.

    last_timestep = -1
    num_time_steps = 0
    num_trajectories = 0
    for ex in generator_utils.tfrecord_iterator(
        filenames, example_spec=ep.example_reading_spec()[0]):
      num_time_steps += 1
      this_timestep = ex[env_problem.TIMESTEP_FIELD][0]
      if 1 + last_timestep != this_timestep:
        num_trajectories += 1
        self.assertEqual(0, this_timestep)
      last_timestep = this_timestep
    num_trajectories += 1

    return num_trajectories, num_time_steps

  def play_env(self,
               env=None,
               nsteps=100,
               base_env_name=None,
               batch_size=5,
               reward_range=None):
    """Creates `EnvProblem` with the given arguments and plays it randomly.

    Args:
      env: optional env.
      nsteps: plays the env randomly for nsteps.
      base_env_name: passed to EnvProblem's init.
      batch_size: passed to EnvProblem's init.
      reward_range: passed to EnvProblem's init.

    Returns:
      tuple of env_problem, number of trajectories done, number of trajectories
      done in the last step.
    """

    if env is None:
      env = env_problem.EnvProblem(
          base_env_name=base_env_name,
          batch_size=batch_size,
          reward_range=reward_range)
      # Usually done by a registered subclass, we do this manually in the test.
      env.name = base_env_name

    # Reset all environments.
    env.reset()

    # Play for some steps to generate data.
    num_dones = 0
    num_dones_in_last_step = 0
    for _ in range(nsteps):
      # Sample actions.
      actions = np.stack([env.action_space.sample() for _ in range(batch_size)])
      # Step through it.
      _, _, dones, _ = env.step(actions)
      # Get the indices where we are done ...
      done_indices = env_problem_utils.done_indices(dones)
      # ... and reset those.
      env.reset(indices=done_indices)
      # count the number of dones we got, in this step and overall.
      num_dones_in_last_step = sum(dones)
      num_dones += num_dones_in_last_step

    return env, num_dones, num_dones_in_last_step

  def test_generate_data(self):
    base_env_name = "CartPole-v0"
    batch_size = 5
    reward_range = (-1, 1)
    nsteps = 100
    ep, num_dones, num_dones_in_last_step = self.play_env(
        base_env_name=base_env_name,
        batch_size=batch_size,
        reward_range=reward_range,
        nsteps=nsteps)

    # This is because every num_dones starts a new trajectory, and a further
    # batch_size are active at the last step when we call generate_data, but
    # the ones that got done in the last step (these have only one time-step in
    # their trajectory) will be skipped.
    expected_num_trajectories = num_dones + batch_size - num_dones_in_last_step

    # Similar logic as above, nsteps * batch_size overall `step` calls are made.
    expected_num_time_steps = (
        nsteps * batch_size) + num_dones + batch_size - num_dones_in_last_step

    # Dump the completed data to disk.
    ep.generate_data(self.tmp_dir, self.tmp_dir)

    # Read the written files and assert on the number of time steps.
    training_filenames = ep.training_filepaths(
        self.tmp_dir, ep.num_shards[problem.DatasetSplit.TRAIN], True)
    dev_filenames = ep.dev_filepaths(
        self.tmp_dir, ep.num_shards[problem.DatasetSplit.EVAL], True)

    training_trajectories, training_timesteps = self.read_tfrecord_dataset(
        training_filenames, ep)
    dev_trajectories, dev_timesteps = self.read_tfrecord_dataset(
        dev_filenames, ep)

    # This tests what we wrote on disk matches with what we computed.
    self.assertEqual(expected_num_time_steps,
                     training_timesteps + dev_timesteps)
    self.assertEqual(expected_num_trajectories,
                     training_trajectories + dev_trajectories)

  def test_problem_dataset_works(self):

    # We need to derive this class to set the required methods.
    class TestEnv(env_problem.EnvProblem):
      name = "TestEnv"

      @property
      def input_modality(self):
        return modalities.ModalityType.REAL_L2_LOSS

      @property
      def input_vocab_size(self):
        return None

      @property
      def target_modality(self):
        return modalities.ModalityType.SYMBOL_WEIGHTS_ALL

      @property
      def target_vocab_size(self):
        return 2

    base_env_name = "CartPole-v0"
    batch_size = 5
    reward_range = (-1, 1)

    env = TestEnv(
        base_env_name=base_env_name,
        batch_size=batch_size,
        reward_range=reward_range)

    nsteps = 100
    ep, _, _ = self.play_env(env=env, nsteps=nsteps)

    # Dump the completed data to disk.
    ep.generate_data(self.tmp_dir, self.tmp_dir)

    # Read the actual files and count the trajectories and time-steps.
    dev_filenames = ep.dev_filepaths(
        self.tmp_dir, ep.num_shards[problem.DatasetSplit.EVAL], True)
    dev_trajectories, dev_timesteps = self.read_tfrecord_dataset(
        dev_filenames, ep)

    # Count them using a tf.data.Dataset.
    dev_dataset = ep.dataset(tf.estimator.ModeKeys.EVAL, data_dir=self.tmp_dir)

    last_timestep = -1
    dev_timesteps_ds = 0
    dev_trajectories_ds = 0
    iterator = dev_dataset.make_one_shot_iterator()
    next_element = iterator.get_next()
    with tf.Session() as session:
      while True:
        try:
          tf_example_dict = session.run(next_element)

          # We have a time-step.
          dev_timesteps_ds += 1

          this_timestep = tf_example_dict[env_problem.TIMESTEP_FIELD][
              0]  # [0] since every value in tf_example_dict is an array/list.
          if 1 + last_timestep != this_timestep:
            dev_trajectories_ds += 1
            self.assertEqual(0, this_timestep)
          last_timestep = this_timestep
        except tf.errors.OutOfRangeError:
          dev_trajectories_ds += 1
          break

    # Make sure that they agree.
    self.assertEqual(dev_trajectories, dev_trajectories_ds)
    self.assertEqual(dev_timesteps, dev_timesteps_ds)

if __name__ == "__main__":
  tf.test.main()
<commit_msg>Work around differences in internal and external Gym.<commit_after># coding=utf-8
# Copyright 2019 The Tensor2Tensor Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for tensor2tensor.envs.env_problem."""

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

import os
from gym.spaces import Box
from gym.spaces import Discrete
import numpy as np
from tensor2tensor.data_generators import generator_utils
from tensor2tensor.data_generators import problem
from tensor2tensor.envs import env_problem
from tensor2tensor.envs import env_problem_utils
from tensor2tensor.layers import modalities
import tensorflow as tf

class EnvProblemTest(tf.test.TestCase):

  def setUp(self):
    self.tmp_dir = os.path.join(tf.test.get_temp_dir(), "tmp_dir")
    tf.gfile.MakeDirs(self.tmp_dir)

  def tearDown(self):
    tf.gfile.DeleteRecursively(self.tmp_dir)

  def test_setup(self):
    ep = env_problem.EnvProblem(base_env_name="CartPole-v0", batch_size=5)
    # Checks that environments were created and they are `batch_size` in number.
    ep.assert_common_preconditions()

    # Expectations on the observation space.
    observation_space = ep.observation_space
    self.assertTrue(isinstance(observation_space, Box))
    self.assertEqual(observation_space.shape, (4,))
    self.assertEqual(observation_space.dtype, np.float32)

    # Expectations on the action space.
    action_space = ep.action_space
    self.assertTrue(isinstance(action_space, Discrete))
    self.assertEqual(action_space.shape, ())
    self.assertEqual(action_space.dtype, np.int64)
    self.assertEqual(ep.num_actions, 2)

    # Reward range is infinite here.
    self.assertFalse(ep.is_reward_range_finite)

  def test_reward_range(self):
    # Passing reward_range=None means take the reward range of the underlying
    # environment as the reward range.
    ep = env_problem.EnvProblem(
        base_env_name="FrozenLake-v0", batch_size=5, reward_range=None)
    ep.assert_common_preconditions()

    # Assert reward range is finite here.
    self.assertTrue(ep.is_reward_range_finite)

    # Assert that it is as expected of the underlying environment, since reward_
    self.assertEqual(0, ep.reward_range[0])
    self.assertEqual(1, ep.reward_range[1])

  def test_default_processed_rewards_discrete(self):
    # This differs in the above because it has a Tuple observation space.
    ep = env_problem.EnvProblem(
        base_env_name="KellyCoinflip-v0", batch_size=5, reward_range=None)
    ep.assert_common_preconditions()

    # Assert reward range is finite here.
    self.assertTrue(ep.is_reward_range_finite)

    # Assert that it is as expected of the underlying environment.
    reward_range = ep.reward_range
    self.assertEqual(0, reward_range[0])

    # Google's version of Gym has maxWealth, vs max_wealth externally.
    max_wealth = getattr(ep._envs[0], "maxWealth",
                         getattr(ep._envs[0], "max_wealth", None))
    self.assertIsNotNone(max_wealth)
    self.assertEqual(max_wealth, reward_range[1])

    # Check that the processed rewards are discrete.
    self.assertTrue(ep.is_processed_rewards_discrete)

    # Assert on the number of rewards.
    self.assertEqual(ep.num_rewards, reward_range[1] - reward_range[0] + 1)

  def test_interaction_with_env(self):
    batch_size = 5
    reward_range = (-1, 1)
    ep = env_problem.EnvProblem(
        base_env_name="KellyCoinflip-v0",
        batch_size=batch_size,
        reward_range=reward_range)

    # Resets all environments.
    ep.reset()

    # Let's play a few steps.
    nsteps = 100
    num_trajectories_completed = 0
    num_timesteps_completed = 0
    # If batch_done_at_step[i] = j then it means that i^th env last got done at
    # step = j.
    batch_done_at_step = np.full(batch_size, -1)
    for i in range(nsteps):
      # Sample batch_size actions from the action space and stack them (since
      # that is the expected type).
      actions = np.stack([ep.action_space.sample() for _ in range(batch_size)])

      _, _, dones, _ = ep.step(actions)

      # Do the book-keeping on number of trajectories completed and expect that
      # it matches ep's completed number.

      num_done = sum(dones)
      num_trajectories_completed += num_done

      self.assertEqual(num_trajectories_completed,
                       len(ep.trajectories.completed_trajectories))

      # Get the indices where we are done ...
      done_indices = env_problem_utils.done_indices(dones)

      # ... and reset those.
      ep.reset(indices=done_indices)

      # If nothing got done, go on to the next step.
      if done_indices.size == 0:
        # i.e. this is an empty array.
        continue

      # See when these indices were last done and calculate how many time-steps
      # each one took to get done.
      num_timesteps_completed += sum(i + 1 - batch_done_at_step[done_indices])
      batch_done_at_step[done_indices] = i

      # This should also match the number of time-steps completed given by ep.
      num_timesteps_completed_ep = sum(
          ct.num_time_steps for ct in ep.trajectories.completed_trajectories)
      self.assertEqual(num_timesteps_completed, num_timesteps_completed_ep)

    # Reset the trajectories.
    ep.trajectories.reset_batch_trajectories()
    self.assertEqual(0, len(ep.trajectories.completed_trajectories))

  def read_tfrecord_dataset(self, filenames, ep):
    # Read the dataset at `filenames` into a tf.data.Dataset and returns the
    # number of time-steps (just the number of records in the dataset) and the
    # number of trajectories.

    last_timestep = -1
    num_time_steps = 0
    num_trajectories = 0
    for ex in generator_utils.tfrecord_iterator(
        filenames, example_spec=ep.example_reading_spec()[0]):
      num_time_steps += 1
      this_timestep = ex[env_problem.TIMESTEP_FIELD][0]
      if 1 + last_timestep != this_timestep:
        num_trajectories += 1
        self.assertEqual(0, this_timestep)
      last_timestep = this_timestep
    num_trajectories += 1

    return num_trajectories, num_time_steps

  def play_env(self,
               env=None,
               nsteps=100,
               base_env_name=None,
               batch_size=5,
               reward_range=None):
    """Creates `EnvProblem` with the given arguments and plays it randomly.

    Args:
      env: optional env.
      nsteps: plays the env randomly for nsteps.
      base_env_name: passed to EnvProblem's init.
      batch_size: passed to EnvProblem's init.
      reward_range: passed to EnvProblem's init.

    Returns:
      tuple of env_problem, number of trajectories done, number of trajectories
      done in the last step.
    """

    if env is None:
      env = env_problem.EnvProblem(
          base_env_name=base_env_name,
          batch_size=batch_size,
          reward_range=reward_range)
      # Usually done by a registered subclass, we do this manually in the test.
      env.name = base_env_name

    # Reset all environments.
    env.reset()

    # Play for some steps to generate data.
    num_dones = 0
    num_dones_in_last_step = 0
    for _ in range(nsteps):
      # Sample actions.
      actions = np.stack([env.action_space.sample() for _ in range(batch_size)])
      # Step through it.
      _, _, dones, _ = env.step(actions)
      # Get the indices where we are done ...
      done_indices = env_problem_utils.done_indices(dones)
      # ... and reset those.
      env.reset(indices=done_indices)
      # count the number of dones we got, in this step and overall.
      num_dones_in_last_step = sum(dones)
      num_dones += num_dones_in_last_step

    return env, num_dones, num_dones_in_last_step

  def test_generate_data(self):
    base_env_name = "CartPole-v0"
    batch_size = 5
    reward_range = (-1, 1)
    nsteps = 100
    ep, num_dones, num_dones_in_last_step = self.play_env(
        base_env_name=base_env_name,
        batch_size=batch_size,
        reward_range=reward_range,
        nsteps=nsteps)

    # This is because every num_dones starts a new trajectory, and a further
    # batch_size are active at the last step when we call generate_data, but
    # the ones that got done in the last step (these have only one time-step in
    # their trajectory) will be skipped.
    expected_num_trajectories = num_dones + batch_size - num_dones_in_last_step

    # Similar logic as above, nsteps * batch_size overall `step` calls are made.
    expected_num_time_steps = (
        nsteps * batch_size) + num_dones + batch_size - num_dones_in_last_step

    # Dump the completed data to disk.
    ep.generate_data(self.tmp_dir, self.tmp_dir)

    # Read the written files and assert on the number of time steps.
    training_filenames = ep.training_filepaths(
        self.tmp_dir, ep.num_shards[problem.DatasetSplit.TRAIN], True)
    dev_filenames = ep.dev_filepaths(
        self.tmp_dir, ep.num_shards[problem.DatasetSplit.EVAL], True)

    training_trajectories, training_timesteps = self.read_tfrecord_dataset(
        training_filenames, ep)
    dev_trajectories, dev_timesteps = self.read_tfrecord_dataset(
        dev_filenames, ep)

    # This tests what we wrote on disk matches with what we computed.
    self.assertEqual(expected_num_time_steps,
                     training_timesteps + dev_timesteps)
    self.assertEqual(expected_num_trajectories,
                     training_trajectories + dev_trajectories)

  def test_problem_dataset_works(self):

    # We need to derive this class to set the required methods.
    class TestEnv(env_problem.EnvProblem):
      name = "TestEnv"

      @property
      def input_modality(self):
        return modalities.ModalityType.REAL_L2_LOSS

      @property
      def input_vocab_size(self):
        return None

      @property
      def target_modality(self):
        return modalities.ModalityType.SYMBOL_WEIGHTS_ALL

      @property
      def target_vocab_size(self):
        return 2

    base_env_name = "CartPole-v0"
    batch_size = 5
    reward_range = (-1, 1)

    env = TestEnv(
        base_env_name=base_env_name,
        batch_size=batch_size,
        reward_range=reward_range)

    nsteps = 100
    ep, _, _ = self.play_env(env=env, nsteps=nsteps)

    # Dump the completed data to disk.
    ep.generate_data(self.tmp_dir, self.tmp_dir)

    # Read the actual files and count the trajectories and time-steps.
    dev_filenames = ep.dev_filepaths(
        self.tmp_dir, ep.num_shards[problem.DatasetSplit.EVAL], True)
    dev_trajectories, dev_timesteps = self.read_tfrecord_dataset(
        dev_filenames, ep)

    # Count them using a tf.data.Dataset.
    dev_dataset = ep.dataset(tf.estimator.ModeKeys.EVAL, data_dir=self.tmp_dir)

    last_timestep = -1
    dev_timesteps_ds = 0
    dev_trajectories_ds = 0
    iterator = dev_dataset.make_one_shot_iterator()
    next_element = iterator.get_next()
    with tf.Session() as session:
      while True:
        try:
          tf_example_dict = session.run(next_element)

          # We have a time-step.
          dev_timesteps_ds += 1

          this_timestep = tf_example_dict[env_problem.TIMESTEP_FIELD][
              0]  # [0] since every value in tf_example_dict is an array/list.
          if 1 + last_timestep != this_timestep:
            dev_trajectories_ds += 1
            self.assertEqual(0, this_timestep)
          last_timestep = this_timestep
        except tf.errors.OutOfRangeError:
          dev_trajectories_ds += 1
          break

    # Make sure that they agree.
    self.assertEqual(dev_trajectories, dev_trajectories_ds)
    self.assertEqual(dev_timesteps, dev_timesteps_ds)

if __name__ == "__main__":
  tf.test.main()
<|endoftext|>


<commit_before>from datetime import datetime

from djpcms import forms, views, html
from djpcms.template import loader
from djpcms.apps.included.admin import AdminApplication, AdminApplicationSimple
from djpcms.utils.dates import nicetimedelta
from djpcms.utils.text import nicename
from djpcms.utils import mark_safe
from djpcms.forms.utils import return_form_errors

from stdnet import orm

import pulsar
from pulsar.utils.py2py3 import iteritems
from pulsar.apps import tasks
from pulsar.http import rpc

from .models import Task

fromtimestamp = datetime.fromtimestamp

class ServerForm(forms.Form):
    code = forms.CharField()
    schema = forms.CharField(initial = 'http://')
    host = forms.CharField()
    port = forms.IntegerField(initial = 8060)
    notes = forms.CharField(widget = html.TextArea,
                            required = False)
    location = forms.CharField(required = False)

class PulsarServerApplication(AdminApplication):
    inherit = True
    form = ServerForm
    list_per_page = 100
    template_view = ('pulsardjp/monitor.html',)
    converters = {'uptime': nicetimedelta}
    list_display = ('code','path','machine','this','notes')

    def get_client(self, instance):
        return rpc.JsonProxy(instance.path())

    def render_object_view(self, djp):
        r = self.get_client(djp.instance)
        try:
            panels = self.get_panels(djp,r.server_info())
        except pulsar.ConnectionError:
            panels = {'left_panels':[{'name':'Server','value':'No Connection'}]}
        return loader.render(self.template_view,panels)

    def pannel_data(self, data):
        for k,v in iteritems(data):
            if k in self.converters:
                v = self.converters[k](v)
            yield {'name':nicename(k),
                   'value':v}

    def get_panels(self,djp,info):
        monitors = []
        for monitor in info['monitors']:
            monitors.append({'name':nicename(monitor.pop('name','Monitor')),
                             'value':html.ObjectDefinition(self,djp,\
                                          self.pannel_data(monitor))})
        servers = [{'name':'Server',
                    'value':html.ObjectDefinition(self,djp,\
                                   self.pannel_data(info['server']))}]
        return {'left_panels':servers,
                'right_panels':monitors}

task_display = ('name','status','timeout','time_executed',
                'time_start','time_end','duration','expiry',
                'user')

class TasksAdmin(AdminApplicationSimple):
    list_display = ('short_id',) + task_display
    object_display = ('id',) + task_display + ('api','string_result','stack_trace')
    has_plugins = False
    inherit = False
    search = views.SearchView()
    view   = views.ViewView(regex = views.UUID_REGEX)
    delete = views.DeleteView()

#
# Scripts
#

script_languages = (
                    ('python','python'),
                    )

class ScriptForm(forms.Form):
    name = forms.CharField(toslug = '_')
    language = forms.ChoiceField(choices = script_languages)
    body = forms.CharField(widget = html.TextArea(default_class = 'taboverride'))

    def clean_name(self, value):
        return orm.test_unique('name',self.model,value,self.instance,
                               forms.ValidationError)

class RunScriptForm(forms.Form):
    parameters = forms.CharField(required = False)

    def clean_parameters(self, value):
        return value

class RunScriptView(views.ChangeView):

    def default_post(self, djp):
        fhtml = self.get_form(djp)
        form = fhtml.form
        if form.is_valid():
            self.appmodel.run(djp)
        else:
            return return_form_errors(fhtml,djp)

HtmlRunScriptForm = forms.HtmlForm(
    RunScriptForm,
    #layout = Layout(default_style = blockLabels2),
    inputs = (('run','_save'),)
)

class ScriptApplication(views.ModelApplication):
    inherit = True
    form = ScriptForm
    list_display = ('name','language','parameters')

    run_view = RunScriptView(regex = 'run', form = HtmlRunScriptForm)

    class Media:
        js = ['djpcms/taboverride.js']

    def run(self, djp, **params):
        '''This needs to be implemented by your application'''
        pass

<commit_msg>minor<commit_after>from datetime import datetime

from djpcms import forms, views, html
from djpcms.template import loader
from djpcms.apps.included.admin import AdminApplication, AdminApplicationSimple
from djpcms.utils.dates import nicetimedelta
from djpcms.utils.text import nicename
from djpcms.utils import mark_safe
from djpcms.forms.utils import return_form_errors

from stdnet import orm

import pulsar
from pulsar.utils.py2py3 import iteritems
from pulsar.apps import tasks
from pulsar.http import rpc

from .models import Task

fromtimestamp = datetime.fromtimestamp

class ServerForm(forms.Form):
    code = forms.CharField()
    schema = forms.CharField(initial = 'http://')
    host = forms.CharField()
    port = forms.IntegerField(initial = 8060)
    notes = forms.CharField(widget = html.TextArea,
                            required = False)
    location = forms.CharField(required = False)

class PulsarServerApplication(AdminApplication):
    inherit = True
    form = ServerForm
    list_per_page = 100
    template_view = ('pulsardjp/monitor.html',)
    converters = {'uptime': nicetimedelta}
    list_display = ('code','path','machine','this','notes')

    def get_client(self, instance):
        return rpc.JsonProxy(instance.path())

    def render_object_view(self, djp):
        r = self.get_client(djp.instance)
        try:
            panels = self.get_panels(djp,r.server_info())
        except pulsar.ConnectionError:
            panels = {'left_panels':[{'name':'Server','value':'No Connection'}]}
        return loader.render(self.template_view,panels)

    def pannel_data(self, data):
        for k,v in iteritems(data):
            if k in self.converters:
                v = self.converters[k](v)
            yield {'name':nicename(k),
                   'value':v}

    def get_panels(self,djp,info):
        monitors = []
        for monitor in info['monitors']:
            monitors.append({'name':nicename(monitor.pop('name','Monitor')),
                             'value':html.ObjectDefinition(self,djp,\
                                          self.pannel_data(monitor))})
        servers = [{'name':'Server',
                    'value':html.ObjectDefinition(self,djp,\
                                   self.pannel_data(info['server']))}]
        return {'left_panels':servers,
                'right_panels':monitors}

task_display = ('name','status','timeout','time_executed',
                'time_start','time_end','duration','expiry',
                'user')

class TasksAdmin(AdminApplicationSimple):
    list_display = ('short_id',) + task_display
    object_display = ('id',) + task_display + ('api','string_result','stack_trace')
    has_plugins = False
    inherit = False
    search = views.SearchView()
    view   = views.ViewView(regex = views.UUID_REGEX)
    delete = views.DeleteView()

#
# Scripts
#

script_languages = (
                    ('python','python'),
                    )

class ScriptForm(forms.Form):
    name = forms.CharField(toslug = '_')
    language = forms.ChoiceField(choices = script_languages)
    body = forms.CharField(widget = html.TextArea(default_class = 'taboverride'))

    def clean_name(self, value):
        return orm.test_unique('name',self.model,value,self.instance,
                               forms.ValidationError)

class RunScriptForm(forms.Form):
    parameters = forms.CharField(required = False)

    def clean_parameters(self, value):
        return value

class RunScriptView(views.ChangeView):

    def default_post(self, djp):
        fhtml = self.get_form(djp)
        form = fhtml.form
        if form.is_valid():
            self.appmodel.run(djp)
        else:
            return return_form_errors(fhtml,djp)

HtmlRunScriptForm = forms.HtmlForm(
    RunScriptForm,
    #layout = Layout(default_style = blockLabels2),
    inputs = (('run','_save'),)
)

class ScriptApplication(views.ModelApplication):
    inherit = True
    form = ScriptForm
    list_display = ('name','language','parameters')

    run_view = RunScriptView(regex = 'run', form = HtmlRunScriptForm)

    class Media:
        js = ['djpcms/taboverride.js']

    def run(self, djp, **params):
        '''This needs to be implemented by your application'''
        pass

<|endoftext|>


<commit_before>#!/usr/bin/env python
import os
import sys
import pandas as pd
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import argparse

def get_args():
    """Gets the command line arguments"""
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument(
        'results_file',
        help='The input results spreadsheet')
    parser.add_argument(
        '--worksheet_name',
        type=str,
        default='results',
        help='Worksheet containing the results to plot')
    parser.add_argument(
        '--compute_element_name',
        type=str,
        default='Threads',
        help='Compute element name that will be used on plot labels')
    parser.add_argument(
        '--walltime_units',
        default='Minutes',
        type=str,
        help='Walltime units name')
    parser.add_argument(
        '--filter_column',
        default='',
        type=str,
        help='''Optional filter column name. If given, the column contents will
        be equated to a boolean value. Any result group with a member filter
        that equates to False will be excluded from the plot.''')
    parser.add_argument(
        '--title_prefix',
        type=str,
        help='optional prefix that will be prepended to the standard plot titles')
    parser.add_argument(
        '--file_prefix',
        type=str,
        help='optional prefix that will be prepended to the standard file names')
    parser.add_argument(
        '--window',
        default=False,
        help='prints plots to a window instead of to files',
        action='store_true')
    parser.add_argument(
        '--plot_width',
        default=10,
        type=int,
        help='Plot width in inches')
    parser.add_argument(
        '--speedup_max',
        default=16.1,
        type=float,
        help='Max Y axis scale for speedup')
    parser.add_argument(
        '--weak',
        default=False,
        help='Generate plots for weak scaling instead of strong scaling',
        action='store_true')
    parser.add_argument(
        '--style',
        default='',
        type=str,
        help='If available, use one of the matplotlib predefined styles')

    return parser.parse_args()

def save(path, ext='png', close=True, verbose=True):
    """Save a figure from pyplot.

    Parameters
    ----------
    path : string
        The path (and filename, without the extension) to save the
        figure to.

    ext : string (default='png')
        The file extension. This must be supported by the active
        matplotlib backend (see matplotlib.backends module).  Most
        backends support 'png', 'pdf', 'ps', 'eps', and 'svg'.

    close : boolean (default=True)
        Whether to close the figure after saving.  If you want to save
        the figure multiple times (e.g., to multiple formats), you
        should NOT close it in between saves or you will have to
        re-plot it.

    verbose : boolean (default=True)
        Whether to print information about when and where the image
        has been saved.

    """

    # Extract the directory and filename from the given path
    directory = os.path.split(path)[0]
    filename = "%s.%s" % (os.path.split(path)[1], ext)
    if directory == '':
        directory = '.'

    # If the directory does not exist, create it
    if not os.path.exists(directory):
        os.makedirs(directory)

    # The final path to save to
    savepath = os.path.join(directory, filename)

    if verbose:
        print("Saving figure to '%s'..." % savepath),

    # Actually save the figure
    plt.savefig(savepath)

    # Close it
    if close:
        plt.close()

    if verbose:
        print("Done")

def read_dataframe_from_excel(
        filename,
        worksheet,
        **kwargs):
    """Reads a pandas dataframe from an Excel file
    filename: the excel file
    worksheet: worksheet name within the excel file
    The remaining named arguments are passed to pandas.read_excel
    """
    results = pd.read_excel(filename, worksheet, **kwargs)
    return results

def calculate_speedup_and_efficiency(
        rdf,
        compute_element_col_index,
        time_col_index):
    """Calculates the speedup and efficiency and
    adds them as new columns to a dataframe

    rdf: the results dataframe
    compute_element_col_index: compute elements column name
    time_col_index: column name containing the computation times

    Returns: a tuple containing the speedup, strong scaling efficiency, and
    weak scaling efficiency series
    """
    # get the t1 reference value
    t1 = float(rdf.loc[rdf[compute_element_col_index] == 1][time_col_index])

    # speedup n = t1 / tn
    speedup = t1 / rdf[time_col_index]

    # strong scaling efficiency n = t1 / (tn * n)
    strong_efficiency = t1 / (rdf[time_col_index] * rdf[compute_element_col_index])

    # weak scaling efficiency n = t1 / tn
    weak_efficiency = t1 / rdf[time_col_index]

    return (speedup, strong_efficiency, weak_efficiency)

def plot_walltime(
        series,
        colours,
        series_names,
        compute_elements,
        ymax,
        plot_size=None,
        group_width=0.8,
        xlabel='Processes',
        ylabel='Minutes',
        title='Walltime',
        file_name='walltime',
        file_extension='png',
        y_log_scale=False,
        show=False):
    """creates a bar plot as a new pyplot figure"""

    # the x locations for the groups.
    x_ind = np.arange(len(compute_elements))

    # create the figure and axes
    fig = plt.figure(figsize=plot_size)
    ax = fig.add_subplot(111)

    # create the individual bars
    bars_per_group = len(series)
    bar_width = group_width / bars_per_group
    bar_left = x_ind - (group_width / 2)  # centre the group on the ticks
    plot_left = x_ind[0] - group_width
    plot_right = x_ind[-1] + group_width

    for values, colour, name in zip(series, colours, series_names):
        ax.bar(
            bar_left,
            values,
            width=bar_width,
            color=colour,
            label=name,
            log=y_log_scale)
        bar_left += bar_width

    # apply the labels and formatting
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.set_xlim(plot_left, plot_right)
    ax.set_ylim(0, ymax)
    ax.set_xticks(x_ind)
    ax.set_xticklabels(compute_elements)
    ax.legend()
    ax.get_legend().get_frame().set_facecolor('white')

    if show:
        plt.show()
    else:
        save(file_name, file_extension, close=True, verbose=True)

def plot_speedup(
        series,
        colours,
        series_names,
        compute_elements,
        xmax,
        ymax,
        line_width=1,
        plot_size=None,
        xlabel='Processes',
        ylabel='Speedup',
        title='Speedup',
        file_name='speedup',
        file_extension='png',
        show=False):
    """creates a speedup plot"""

    # create the figure and axes
    fig = plt.figure(figsize=plot_size)
    ax = fig.add_subplot(111)

    # define the sizes and locations of things
    x_ticks = list(compute_elements)
    x_ticks.append(xmax)
    plot_left = x_ticks[0]
    plot_right = x_ticks[-1]
    ax.set_xlim(plot_left, plot_right)
    ax.set_ylim(plot_left, ymax)

    # create the ideal speedup reference line
    ax.plot(
        x_ticks,
        x_ticks,
        color='red',
        linewidth=line_width,
        label='_nolegend_')

    # plot each series
    for values, colour, name in zip(series, colours, series_names):
        ax.plot(
            compute_elements,
            values,
            label=name,
            color=colour,
            linewidth=line_width,
            marker='o')

    # apply the labels and formatting
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.set_xticks(compute_elements)
    ax.set_xticklabels(compute_elements)
    ax.legend()
    ax.get_legend().get_frame().set_facecolor('white')

    if show:
        plt.show()
    else:
        save(file_name, file_extension, close=True, verbose=True)

def plot_efficiency(
        series,
        colours,
        series_names,
        compute_elements,
        xmax,
        ymax=1.2,
        line_width=1,
        plot_size=None,
        xlabel='Processes',
        ylabel='Efficiency',
        title='Efficiency',
        file_name='efficiency',
        file_extension='png',
        show=False):
    """creates an efficiency plot"""

    # create the figure and axes
    fig = plt.figure(figsize=plot_size)
    ax = fig.add_subplot(111)

    # define the sizes and locations of things
    x_ticks = list(compute_elements)
    x_ticks.append(xmax)
    plot_left = x_ticks[0]
    plot_right = x_ticks[-1]
    ax.set_xlim(plot_left, plot_right)
    ax.set_ylim(0, ymax)

    # create a horizontal ideal efficiency reference line
    ax.hlines(
        1,
        plot_left,
        plot_right,
        colors='red',
        linestyles='solid',
        linewidth=line_width,
        label='_nolegend_')

    # plot each series
    for values, colour, name in zip(series, colours, series_names):
        ax.plot(
            compute_elements,
            values,
            label=name,
            color=colour,
            linewidth=line_width,
            marker='o')

    # apply the labels and formatting
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.set_xticks(compute_elements)
    ax.set_xticklabels(compute_elements)
    ax.legend()
    ax.get_legend().get_frame().set_facecolor('white')

    if show:
        plt.show()
    else:
        save(file_name, file_extension, close=True, verbose=True)

def add_optional_prefix(base_string, prefix, separator):
    result = str(base_string)
    if prefix:
        result = prefix + separator + base_string
    return result

if __name__ == '__main__':

    # get and process the arguments
    args = get_args()

    # If we have a new enough version of matplotlib to support styles, use ggplot
    if args.style:
        try:
            matplotlib.style.use(args.style)
        except:
            pass

    # create plots in a 4:3 aspect ratio
    # matplotlib works in inches
    plot_width = args.plot_width
    plot_size = (plot_width, plot_width * 3 / 4)

    show_instead_of_save = args.window

    walltime_title = add_optional_prefix('Walltime', args.title_prefix, ' - ')
    speedup_title = add_optional_prefix('Speedup', args.title_prefix, ' - ')

    efficiency_title = add_optional_prefix(
        'Efficiency',
        args.title_prefix,
        ' - ')

    walltime_file = add_optional_prefix('walltime', args.file_prefix, '-')
    speedup_file = add_optional_prefix('speedup', args.file_prefix, '-')
    efficiency_file = add_optional_prefix(
        'weak-efficiency' if args.weak else 'strong-efficiency',
        args.file_prefix,
        '-')

    compute_element_name = args.compute_element_name
    walltime_units = args.walltime_units

    # CSIRO colours
    colours = [
        '#00a9ce',  # midday blue
        '#78be20',  # light forest
        '#DF1995',  # fuschia
        '#E87722',  # orange
        '#E4002B',  # vermillion
        '#00616c',  # midnight blue
        '#FFB81C',  # gold
        '#6D2077',  # plum
        '#1E22AA',  # blueberry
    ]

    usecols = [
        'group',
        'compute_elements',
        'walltime']
    if args.filter_column:
        usecols.append(args.filter_column)

    results = read_dataframe_from_excel(
        args.results_file,
        worksheet=args.worksheet_name,
        usecols=usecols)

    # filter out incomplete results
    results = results[results.group.notnull() & (results.walltime > 0)]

    # apply the optional filter column
    if args.filter_column:
        results = results[results[args.filter_column] > 0]

    # if there are multiple times for each (group,compute_element) tuple,
    # calculate the mean
    results = results.groupby(['group','compute_elements']).mean().reset_index()

    # create new dataframes for each group of results
    groups = results.group.unique()
    group_dataframes = {}
    for group in groups:
        df = results[(results.group == group)]
        if len(df):
            group_dataframes[group] = df.sort(
                'compute_elements',
                ascending=True)

    # calculate speedup and efficiency
    for data in group_dataframes.values():
        data['speedup'], data['strong_efficiency'], data['weak_efficiency'] = \
            calculate_speedup_and_efficiency(
                data,
                'compute_elements',
                'walltime')

    # extract the compute element names for grouping and labelling the charts
    compute_elements = np.sort(results.compute_elements.unique())

    # build the collections of values for the charts
    walltimes = []
    strong_efficiencies = []
    weak_efficiencies = []
    speedups = []
    series_names = []
    max_walltime = 0
    for name, data in group_dataframes.items():
        walltimes.append(data.walltime)
        strong_efficiencies.append(data.strong_efficiency)
        weak_efficiencies.append(data.weak_efficiency)
        speedups.append(data.speedup)
        series_names.append(name)
        local_max_walltime = max(data.walltime)
        if local_max_walltime > max_walltime:
            max_walltime = local_max_walltime

    # finally make the plots

    plot_walltime(
        walltimes,
        colours,
        series_names,
        compute_elements,
        ymax=max_walltime * 1.2,
        group_width=0.83,
        plot_size=plot_size,
        title=walltime_title,
        xlabel=compute_element_name,
        ylabel=walltime_units,
        show=show_instead_of_save,
        file_name=walltime_file,
        file_extension='png')

    plot_efficiency(
        weak_efficiencies if args.weak else strong_efficiencies,
        colours,
        series_names,
        compute_elements,
        line_width=1.5,
        xmax=compute_elements.max() * 1.1,
        ymax=1.3,
        plot_size=plot_size,
        xlabel=compute_element_name,
        title=efficiency_title,
        show=show_instead_of_save,
        file_name=efficiency_file,
        file_extension='png')

    # The Speedup plot makes no sense for weak scaling.
    # I could create a new plot for weak scaling that shows the percentage
    # increase in walltime as the compute elements increase
    if not args.weak:
        plot_speedup(
            speedups,
            colours,
            series_names,
            compute_elements,
            line_width=1.5,
            xmax=compute_elements.max() * 1.05,
            ymax=args.speedup_max,
            plot_size=plot_size,
            xlabel=compute_element_name,
            title=speedup_title,
            show=show_instead_of_save,
            file_name=speedup_file,
            file_extension='png')
<commit_msg>made the file type an option<commit_after>#!/usr/bin/env python
import os
import sys
import pandas as pd
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import argparse

def get_args():
    """Gets the command line arguments"""
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument(
        'results_file',
        help='The input results spreadsheet')
    parser.add_argument(
        '--worksheet_name',
        type=str,
        default='results',
        help='Worksheet containing the results to plot')
    parser.add_argument(
        '--compute_element_name',
        type=str,
        default='Threads',
        help='Compute element name that will be used on plot labels')
    parser.add_argument(
        '--walltime_units',
        default='Minutes',
        type=str,
        help='Walltime units name')
    parser.add_argument(
        '--filter_column',
        default='',
        type=str,
        help='''Optional filter column name. If given, the column contents will
        be cast to a boolean value. Any result group with a member filter
        that equates to False will be excluded from the plot.''')
    parser.add_argument(
        '--title_prefix',
        type=str,
        help='optional prefix that will be prepended to the standard plot titles')
    parser.add_argument(
        '--file_prefix',
        type=str,
        help='optional prefix that will be prepended to the standard file names')
    parser.add_argument(
        '--file-extension',
        type=str,
        default='png',
        help="""The file extension. This must be supported by the active
        matplotlib backend (see matplotlib.backends module).  Most
        backends support 'png', 'pdf', 'ps', 'eps', and 'svg'.""")
    parser.add_argument(
        '--window',
        default=False,
        help='prints plots to a window instead of to files',
        action='store_true')
    parser.add_argument(
        '--plot_width',
        default=10,
        type=int,
        help='Plot width in inches')
    parser.add_argument(
        '--speedup_max',
        default=16.1,
        type=float,
        help='Max Y axis scale for speedup')
    parser.add_argument(
        '--weak',
        default=False,
        help='Generate plots for weak scaling instead of strong scaling',
        action='store_true')
    parser.add_argument(
        '--style',
        default='',
        type=str,
        help='If available, use one of the matplotlib predefined styles')

    return parser.parse_args()

def save(path, ext='png', close=True, verbose=True):
    """Save a figure from pyplot.

    Parameters
    ----------
    path : string
        The path (and filename, without the extension) to save the
        figure to.

    ext : string (default='png')
        The file extension. This must be supported by the active
        matplotlib backend (see matplotlib.backends module).  Most
        backends support 'png', 'pdf', 'ps', 'eps', and 'svg'.

    close : boolean (default=True)
        Whether to close the figure after saving.  If you want to save
        the figure multiple times (e.g., to multiple formats), you
        should NOT close it in between saves or you will have to
        re-plot it.

    verbose : boolean (default=True)
        Whether to print information about when and where the image
        has been saved.

    """

    # Extract the directory and filename from the given path
    directory = os.path.split(path)[0]
    filename = "%s.%s" % (os.path.split(path)[1], ext)
    if directory == '':
        directory = '.'

    # If the directory does not exist, create it
    if not os.path.exists(directory):
        os.makedirs(directory)

    # The final path to save to
    savepath = os.path.join(directory, filename)

    if verbose:
        print("Saving figure to '%s'..." % savepath),

    # Actually save the figure
    plt.savefig(savepath)

    # Close it
    if close:
        plt.close()

    if verbose:
        print("Done")

def read_dataframe_from_excel(
        filename,
        worksheet,
        **kwargs):
    """Reads a pandas dataframe from an Excel file
    filename: the excel file
    worksheet: worksheet name within the excel file
    The remaining named arguments are passed to pandas.read_excel
    """
    results = pd.read_excel(filename, worksheet, **kwargs)
    return results

def calculate_speedup_and_efficiency(
        rdf,
        compute_element_col_index,
        time_col_index):
    """Calculates the speedup and efficiency and
    adds them as new columns to a dataframe

    rdf: the results dataframe
    compute_element_col_index: compute elements column name
    time_col_index: column name containing the computation times

    Returns: a tuple containing the speedup, strong scaling efficiency, and
    weak scaling efficiency series
    """
    # get the t1 reference value
    t1 = float(rdf.loc[rdf[compute_element_col_index] == 1][time_col_index])

    # speedup n = t1 / tn
    speedup = t1 / rdf[time_col_index]

    # strong scaling efficiency n = t1 / (tn * n)
    strong_efficiency = t1 / (rdf[time_col_index] * rdf[compute_element_col_index])

    # weak scaling efficiency n = t1 / tn
    weak_efficiency = t1 / rdf[time_col_index]

    return (speedup, strong_efficiency, weak_efficiency)

def plot_walltime(
        series,
        colours,
        series_names,
        compute_elements,
        ymax,
        plot_size=None,
        group_width=0.8,
        xlabel='Processes',
        ylabel='Minutes',
        title='Walltime',
        file_name='walltime',
        file_extension='png',
        y_log_scale=False,
        show=False):
    """creates a bar plot as a new pyplot figure"""

    # the x locations for the groups.
    x_ind = np.arange(len(compute_elements))

    # create the figure and axes
    fig = plt.figure(figsize=plot_size)
    ax = fig.add_subplot(111)

    # create the individual bars
    bars_per_group = len(series)
    bar_width = group_width / bars_per_group
    bar_left = x_ind - (group_width / 2)  # centre the group on the ticks
    plot_left = x_ind[0] - group_width
    plot_right = x_ind[-1] + group_width

    for values, colour, name in zip(series, colours, series_names):
        ax.bar(
            bar_left,
            values,
            width=bar_width,
            color=colour,
            label=name,
            log=y_log_scale)
        bar_left += bar_width

    # apply the labels and formatting
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.set_xlim(plot_left, plot_right)
    ax.set_ylim(0, ymax)
    ax.set_xticks(x_ind)
    ax.set_xticklabels(compute_elements)
    ax.legend()
    ax.get_legend().get_frame().set_facecolor('white')

    if show:
        plt.show()
    else:
        save(file_name, file_extension, close=True, verbose=True)

def plot_speedup(
        series,
        colours,
        series_names,
        compute_elements,
        xmax,
        ymax,
        line_width=1,
        plot_size=None,
        xlabel='Processes',
        ylabel='Speedup',
        title='Speedup',
        file_name='speedup',
        file_extension='png',
        show=False):
    """creates a speedup plot"""

    # create the figure and axes
    fig = plt.figure(figsize=plot_size)
    ax = fig.add_subplot(111)

    # define the sizes and locations of things
    x_ticks = list(compute_elements)
    x_ticks.append(xmax)
    plot_left = x_ticks[0]
    plot_right = x_ticks[-1]
    ax.set_xlim(plot_left, plot_right)
    ax.set_ylim(plot_left, ymax)

    # create the ideal speedup reference line
    ax.plot(
        x_ticks,
        x_ticks,
        color='red',
        linewidth=line_width,
        label='_nolegend_')

    # plot each series
    for values, colour, name in zip(series, colours, series_names):
        ax.plot(
            compute_elements,
            values,
            label=name,
            color=colour,
            linewidth=line_width,
            marker='o')

    # apply the labels and formatting
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.set_xticks(compute_elements)
    ax.set_xticklabels(compute_elements)
    ax.legend()
    ax.get_legend().get_frame().set_facecolor('white')

    if show:
        plt.show()
    else:
        save(file_name, file_extension, close=True, verbose=True)

def plot_efficiency(
        series,
        colours,
        series_names,
        compute_elements,
        xmax,
        ymax=1.2,
        line_width=1,
        plot_size=None,
        xlabel='Processes',
        ylabel='Efficiency',
        title='Efficiency',
        file_name='efficiency',
        file_extension='png',
        show=False):
    """creates an efficiency plot"""

    # create the figure and axes
    fig = plt.figure(figsize=plot_size)
    ax = fig.add_subplot(111)

    # define the sizes and locations of things
    x_ticks = list(compute_elements)
    x_ticks.append(xmax)
    plot_left = x_ticks[0]
    plot_right = x_ticks[-1]
    ax.set_xlim(plot_left, plot_right)
    ax.set_ylim(0, ymax)

    # create a horizontal ideal efficiency reference line
    ax.hlines(
        1,
        plot_left,
        plot_right,
        colors='red',
        linestyles='solid',
        linewidth=line_width,
        label='_nolegend_')

    # plot each series
    for values, colour, name in zip(series, colours, series_names):
        ax.plot(
            compute_elements,
            values,
            label=name,
            color=colour,
            linewidth=line_width,
            marker='o')

    # apply the labels and formatting
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.set_xticks(compute_elements)
    ax.set_xticklabels(compute_elements)
    ax.legend()
    ax.get_legend().get_frame().set_facecolor('white')

    if show:
        plt.show()
    else:
        save(file_name, file_extension, close=True, verbose=True)

def add_optional_prefix(base_string, prefix, separator):
    result = str(base_string)
    if prefix:
        result = prefix + separator + base_string
    return result

if __name__ == '__main__':

    # get and process the arguments
    args = get_args()

    # If we have a new enough version of matplotlib to support styles, use ggplot
    if args.style:
        try:
            matplotlib.style.use(args.style)
        except:
            pass

    # create plots in a 4:3 aspect ratio
    # matplotlib works in inches
    plot_width = args.plot_width
    plot_size = (plot_width, plot_width * 3 / 4)

    show_instead_of_save = args.window

    walltime_title = add_optional_prefix('Walltime', args.title_prefix, ' - ')
    speedup_title = add_optional_prefix('Speedup', args.title_prefix, ' - ')

    efficiency_title = add_optional_prefix(
        'Efficiency',
        args.title_prefix,
        ' - ')

    walltime_file = add_optional_prefix('walltime', args.file_prefix, '-')
    speedup_file = add_optional_prefix('speedup', args.file_prefix, '-')
    efficiency_file = add_optional_prefix(
        'weak-efficiency' if args.weak else 'strong-efficiency',
        args.file_prefix,
        '-')

    compute_element_name = args.compute_element_name
    walltime_units = args.walltime_units

    # CSIRO colours
    colours = [
        '#00a9ce',  # midday blue
        '#78be20',  # light forest
        '#DF1995',  # fuschia
        '#E87722',  # orange
        '#E4002B',  # vermillion
        '#00616c',  # midnight blue
        '#FFB81C',  # gold
        '#6D2077',  # plum
        '#1E22AA',  # blueberry
    ]

    usecols = [
        'group',
        'compute_elements',
        'walltime']
    if args.filter_column:
        usecols.append(args.filter_column)

    results = read_dataframe_from_excel(
        args.results_file,
        worksheet=args.worksheet_name,
        usecols=usecols)

    # filter out incomplete results
    results = results[results.group.notnull() & (results.walltime > 0)]

    # apply the optional filter column
    if args.filter_column:
        results = results[results[args.filter_column] > 0]

    # if there are multiple times for each (group,compute_element) tuple,
    # calculate the mean
    results = results.groupby(['group','compute_elements']).mean().reset_index()

    # create new dataframes for each group of results, sorted by the compute element
    groups = results.group.unique()
    group_dataframes = {}
    for group in groups:
        df = results[(results.group == group)]
        if len(df):
            group_dataframes[group] = df.sort(
                'compute_elements',
                ascending=True)

    # calculate speedup and efficiency
    for data in group_dataframes.values():
        data['speedup'], data['strong_efficiency'], data['weak_efficiency'] = \
            calculate_speedup_and_efficiency(
                data,
                'compute_elements',
                'walltime')

    # extract the compute element names for grouping and labelling the charts
    compute_elements = np.sort(results.compute_elements.unique())

    # build the collections of values for the charts
    walltimes = []
    strong_efficiencies = []
    weak_efficiencies = []
    speedups = []
    series_names = []
    max_walltime = 0
    for name, data in group_dataframes.items():
        walltimes.append(data.walltime)
        strong_efficiencies.append(data.strong_efficiency)
        weak_efficiencies.append(data.weak_efficiency)
        speedups.append(data.speedup)
        series_names.append(name)
        local_max_walltime = max(data.walltime)
        if local_max_walltime > max_walltime:
            max_walltime = local_max_walltime

    # finally make the plots

    plot_walltime(
        walltimes,
        colours,
        series_names,
        compute_elements,
        ymax=max_walltime * 1.2,
        group_width=0.83,
        plot_size=plot_size,
        title=walltime_title,
        xlabel=compute_element_name,
        ylabel=walltime_units,
        show=show_instead_of_save,
        file_name=walltime_file,
        file_extension=args.file_extension)

    plot_efficiency(
        weak_efficiencies if args.weak else strong_efficiencies,
        colours,
        series_names,
        compute_elements,
        line_width=1.5,
        xmax=compute_elements.max() * 1.1,
        ymax=1.3,
        plot_size=plot_size,
        xlabel=compute_element_name,
        title=efficiency_title,
        show=show_instead_of_save,
        file_name=efficiency_file,
        file_extension=args.file_extension)

    # The Speedup plot makes no sense for weak scaling.
    # I could create a new plot for weak scaling that shows the percentage
    # increase in walltime as the compute elements increase
    if not args.weak:
        plot_speedup(
            speedups,
            colours,
            series_names,
            compute_elements,
            line_width=1.5,
            xmax=compute_elements.max() * 1.05,
            ymax=args.speedup_max,
            plot_size=plot_size,
            xlabel=compute_element_name,
            title=speedup_title,
            show=show_instead_of_save,
            file_name=speedup_file,
            file_extension=args.file_extension)
<|endoftext|>


<commit_before>#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import h2o
import sys
import warnings
from h2o.frame import H2OFrame
from h2o.utils.typechecks import assert_is_type, Enum
from pyspark.ml.util import _jvm
from pyspark.rdd import RDD
from pyspark.sql.dataframe import DataFrame
from pyspark.sql.types import *
from pyspark.sql import SparkSession

from ai.h2o.sparkling.FrameConversions import FrameConversions as fc
from ai.h2o.sparkling.Initializer import Initializer
from ai.h2o.sparkling.H2OConf import H2OConf

def _is_of_simple_type(rdd):
    if not isinstance(rdd, RDD):
        raise ValueError('rdd is not of type pyspark.rdd.RDD')

    # Python 3.6 does not contain type long
    # this code ensures we are compatible with both, python 2.7 and python 3.6
    if sys.version_info > (3,):
        type_checks = (str, int, bool, float)
    else:
        type_checks = (str, int, bool, long, float)

    if not rdd.isEmpty() and isinstance(rdd.first(), type_checks):
        return True
    else:
        return False

def _get_first(rdd):
    if rdd.isEmpty():
        raise ValueError('rdd is empty')

    return rdd.first()

class H2OContext(object):

    def __init__(self):
        """
         This constructor is used just to initialize the environment. It does not start H2OContext.
         To start H2OContext use one of the getOrCreate methods. This constructor is internally used in those methods
        """
        try:
            Initializer.load_sparkling_jar()
        except:
            raise

    def __h2o_connect(h2o_context):
        schema = h2o_context._jhc.h2oContext().getConf().getScheme()
        conf = h2o_context._conf

        kwargs = {}
        kwargs["https"] = schema == "https"
        kwargs["verify_ssl_certificates"] = conf.verifySslCertificates()
        if conf.userName() and conf.password():
            kwargs["auth"] = (conf.userName(), conf.password())
        url = "{}://{}:{}".format(schema, h2o_context._client_ip, h2o_context._client_port)
        if conf.contextPath() is not None:
            url = "{}/{}".format(url, conf.contextPath())
        return h2o.connect(url=url, **kwargs)

    @staticmethod
    def getOrCreate(spark=None, conf=None, **kwargs):
        """
        Get existing or create new H2OContext based on provided H2O configuration. If the conf parameter is set then
        configuration from it is used. Otherwise the configuration properties passed to Sparkling Water are used.
        If the values are not found the default values are used in most of the cases. The default cluster mode
        is internal, ie. spark.ext.h2o.external.cluster.mode=false

        :param spark: Spark Context or Spark Session or H2OConf
        :param conf: H2O configuration as instance of H2OConf
        :param kwargs:  additional parameters which are passed to h2o_connect_hook
        :return:  instance of H2OContext
        """

        if spark is not None and not isinstance(spark, H2OConf):
            warnings.warn("Method getOrCreate with spark argument is deprecated. Please use either just getOrCreate() or if you need"
                          "to pass extra H2OConf, use getOrCreate(conf). The spark argument will be removed in release 3.32.")

        # Workaround for bug in Spark 2.1 as SparkSession created in PySpark is not seen in Java
        # and call SparkSession.builder.getOrCreate on Java side creates a new session, which is not
        # desirable
        activeSession = SparkSession._instantiatedSession
        if activeSession is not None:
            jvm = activeSession.sparkContext._jvm
            jvm.org.apache.spark.sql.SparkSession.setDefaultSession(activeSession._jsparkSession)

        if spark is not None and isinstance(spark, H2OConf):
            selected_conf = spark
        elif conf is not None:
            selected_conf = conf
        else:
            selected_conf = H2OConf()
        if "auth" in kwargs:
            warnings.warn("Providing authentication via auth field on H2OContext is deprecated. "
                          "Please use setUserName and setPassword setters on H2OConf object.")
            selected_conf.setUserName(kwargs["auth"][0])
            selected_conf.setPassword(kwargs["auth"][1])
            del kwargs["auth"]

        h2o_context = H2OContext()

        # Create backing Java H2OContext
        jhc = _jvm().org.apache.spark.h2o.JavaH2OContext.getOrCreate(selected_conf._jconf)
        h2o_context._jhc = jhc
        h2o_context._conf = selected_conf
        h2o_context._client_ip = jhc.h2oLocalClientIp()
        h2o_context._client_port = jhc.h2oLocalClientPort()

        # Create H2O REST API client
        if not h2o_context.__isClientConnected():
            h2o_context.__h2o_connect()

        h2o_context.__setClientConnected()

        print(h2o_context)

        return h2o_context

    def __isStopped(self):
        hc = self._jhc.h2oContext()
        field = hc.getClass().getSuperclass().getDeclaredField("stopped")
        field.setAccessible(True)
        return field.get(hc)

    def __isClientConnected(self):
        hc = self._jhc.h2oContext()
        field = self.__getClientConnectedField()
        return field.get(hc)

    def __setClientConnected(self):
        hc = self._jhc.h2oContext()
        field = self.__getClientConnectedField()
        field.set(hc, True)

    def __getClientConnectedField(self):
        hc = self._jhc.h2oContext()
        field = hc.getClass().getSuperclass().getDeclaredField("clientConnected")
        field.setAccessible(True)
        return field

    def stop(self):
        h2o.connection().close()
        hc = self._jhc.h2oContext()
        scalaStopMethod = getattr(hc, "org$apache$spark$h2o$H2OContext$$stop")
        scalaStopMethod(False, False, False) # stopSpark = False, stopJVM = False, inShutdownHook = False

        if self._conf.get("spark.ext.h2o.rest.api.based.client", "false") == "false":
            sys.exit()

    def download_h2o_logs(self, destination, container = "ZIP"):
        assert_is_type(container, Enum("ZIP", "LOG"))
        return self._jhc.h2oContext().downloadH2OLogs(destination, container)

    def __str__(self):
        if self.__isClientConnected() and not self.__isStopped():
            return self._jhc.toString()
        else:
            return "H2OContext has been stopped or hasn't been created. Call H2OContext.getOrCreate() or " \
                   "H2OContext.getOrCreate(conf) to create a new one."

    def __repr__(self):
        self.show()
        return ""

    def show(self):
        print(self)

    def get_conf(self):
        return self._conf

    def asSparkFrame(self, h2oFrame, copyMetadata=True):
        """
        Transforms given H2OFrame to Spark DataFrame

        Parameters
        ----------
         h2oFrame : H2OFrame
         copyMetadata: Bool = True

        Returns
        -------
         Spark DataFrame
        """

        if isinstance(h2oFrame, H2OFrame):
            frame_id = h2oFrame.frame_id
            jdf = self._jhc.asDataFrame(frame_id, copyMetadata)
            sqlContext = SparkSession.builder.getOrCreate()._wrapped
            df = DataFrame(jdf, sqlContext)
            # Attach h2o_frame to dataframe which forces python not to delete the frame when we leave the scope of this
            # method.
            # Without this, after leaving this method python would garbage collect the frame since it's not used
            # anywhere and spark. when executing any action on this dataframe, will fail since the frame
            # would be missing.
            df._h2o_frame = h2oFrame
            return df

    def as_spark_frame(self, h2o_frame, copy_metadata=True):
        warnings.warn("Method 'as_spark_frame' is deprecated and will be removed in release 3.30. Please use method 'asSparkFrame' instead!")
        return self.asSparkFrame(h2o_frame, copy_metadata)

    def asH2OFrame(self, sparkFrame, h2oFrameName=None, fullCols=-1):
        """
        Transforms given Spark RDD or DataFrame to H2OFrame.

        Parameters
        ----------
          sparkFrame : Spark RDD or DataFrame
          h2oFrameName : Optional name for resulting H2OFrame
          fullCols : number of first n columns which are sent to the client together with the data

        Returns
        -------
          H2OFrame which contains data of original input Spark data structure
        """
        if isinstance(sparkFrame, DataFrame):
            return fc._as_h2o_frame_from_dataframe(self, sparkFrame, h2oFrameName, fullCols)
        elif isinstance(sparkFrame, RDD) and sparkFrame.isEmpty():
            schema = StructType([])
            empty = self._spark_session.createDataFrame(self._spark_session.sparkContext.emptyRDD(), schema)
            return fc._as_h2o_frame_from_dataframe(self, empty, h2oFrameName, fullCols)
        elif isinstance(sparkFrame, RDD):
            # First check if the type T in RDD[T] is one of the python "primitive" types
            # String, Boolean, Int and Double (Python Long is converted to java.lang.BigInteger)
            if _is_of_simple_type(sparkFrame):
                first = _get_first(sparkFrame)
                # Make this code compatible with python 3.6 and python 2.7
                global long
                if sys.version_info > (3,):
                    long = int

                if isinstance(first, str):
                    return fc._as_h2o_frame_from_RDD_String(self, sparkFrame, h2oFrameName, fullCols)
                elif isinstance(first, bool):
                    return fc._as_h2o_frame_from_RDD_Bool(self, sparkFrame, h2oFrameName, fullCols)
                elif (isinstance(sparkFrame.min(), int) and isinstance(sparkFrame.max(), int)) or (isinstance(sparkFrame.min(), long) and isinstance(sparkFrame.max(), long)):
                    if sparkFrame.min() >= _jvm().Integer.MIN_VALUE and sparkFrame.max() <= _jvm().Integer.MAX_VALUE:
                        return fc._as_h2o_frame_from_RDD_Int(self, sparkFrame, h2oFrameName, fullCols)
                    elif sparkFrame.min() >= _jvm().Long.MIN_VALUE and sparkFrame.max() <= _jvm().Long.MAX_VALUE:
                        return fc._as_h2o_frame_from_RDD_Long(self, sparkFrame, h2oFrameName, fullCols)
                    else:
                        raise ValueError('Numbers in RDD Too Big')
                elif isinstance(first, float):
                    return fc._as_h2o_frame_from_RDD_Float(self, sparkFrame, h2oFrameName, fullCols)
            else:
                return fc._as_h2o_frame_from_complex_type(self, sparkFrame, h2oFrameName, fullCols)
        else:
            raise ValueError('The asH2OFrame method expects Spark DataFrame or RDD as the input only!')

    def as_h2o_frame(self, dataframe, framename=None, full_cols=-1):
        warnings.warn("Method 'as_h2o_frame' is deprecated and will be removed in release 3.30. Please use method 'asH2OFrame' instead!")
        return self.asH2OFrame(dataframe, framename, full_cols)
<commit_msg>[SW-1989] Remove kwargs argument from H2oContext.getOrCreate in python and deprecated verify_ssl_certificates arg handling (#1872)<commit_after>#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import h2o
import sys
import warnings
from h2o.frame import H2OFrame
from h2o.utils.typechecks import assert_is_type, Enum
from pyspark.ml.util import _jvm
from pyspark.rdd import RDD
from pyspark.sql.dataframe import DataFrame
from pyspark.sql.types import *
from pyspark.sql import SparkSession

from ai.h2o.sparkling.FrameConversions import FrameConversions as fc
from ai.h2o.sparkling.Initializer import Initializer
from ai.h2o.sparkling.H2OConf import H2OConf

def _is_of_simple_type(rdd):
    if not isinstance(rdd, RDD):
        raise ValueError('rdd is not of type pyspark.rdd.RDD')

    # Python 3.6 does not contain type long
    # this code ensures we are compatible with both, python 2.7 and python 3.6
    if sys.version_info > (3,):
        type_checks = (str, int, bool, float)
    else:
        type_checks = (str, int, bool, long, float)

    if not rdd.isEmpty() and isinstance(rdd.first(), type_checks):
        return True
    else:
        return False

def _get_first(rdd):
    if rdd.isEmpty():
        raise ValueError('rdd is empty')

    return rdd.first()

class H2OContext(object):

    def __init__(self):
        """
         This constructor is used just to initialize the environment. It does not start H2OContext.
         To start H2OContext use one of the getOrCreate methods. This constructor is internally used in those methods
        """
        try:
            Initializer.load_sparkling_jar()
        except:
            raise

    def __h2o_connect(h2o_context):
        schema = h2o_context._jhc.h2oContext().getConf().getScheme()
        conf = h2o_context._conf

        kwargs = {}
        kwargs["https"] = schema == "https"
        kwargs["verify_ssl_certificates"] = conf.verifySslCertificates()
        if conf.userName() and conf.password():
            kwargs["auth"] = (conf.userName(), conf.password())
        url = "{}://{}:{}".format(schema, h2o_context._client_ip, h2o_context._client_port)
        if conf.contextPath() is not None:
            url = "{}/{}".format(url, conf.contextPath())
        return h2o.connect(url=url, **kwargs)

    @staticmethod
    def getOrCreate(spark=None, conf=None):
        """
        Get existing or create new H2OContext based on provided H2O configuration. If the conf parameter is set then
        configuration from it is used. Otherwise the configuration properties passed to Sparkling Water are used.
        If the values are not found the default values are used in most of the cases. The default cluster mode
        is internal, ie. spark.ext.h2o.external.cluster.mode=false

        :param spark: Spark Context or Spark Session or H2OConf
        :param conf: H2O configuration as instance of H2OConf
        :return:  instance of H2OContext
        """

        if spark is not None and not isinstance(spark, H2OConf):
            warnings.warn("Method getOrCreate with spark argument is deprecated. Please use either just getOrCreate() or if you need"
                          "to pass extra H2OConf, use getOrCreate(conf). The spark argument will be removed in release 3.32.")

        # Workaround for bug in Spark 2.1 as SparkSession created in PySpark is not seen in Java
        # and call SparkSession.builder.getOrCreate on Java side creates a new session, which is not
        # desirable
        activeSession = SparkSession._instantiatedSession
        if activeSession is not None:
            jvm = activeSession.sparkContext._jvm
            jvm.org.apache.spark.sql.SparkSession.setDefaultSession(activeSession._jsparkSession)

        if spark is not None and isinstance(spark, H2OConf):
            selected_conf = spark
        elif conf is not None:
            selected_conf = conf
        else:
            selected_conf = H2OConf()

        h2o_context = H2OContext()

        # Create backing Java H2OContext
        jhc = _jvm().org.apache.spark.h2o.JavaH2OContext.getOrCreate(selected_conf._jconf)
        h2o_context._jhc = jhc
        h2o_context._conf = selected_conf
        h2o_context._client_ip = jhc.h2oLocalClientIp()
        h2o_context._client_port = jhc.h2oLocalClientPort()

        # Create H2O REST API client
        if not h2o_context.__isClientConnected():
            h2o_context.__h2o_connect()

        h2o_context.__setClientConnected()

        print(h2o_context)

        return h2o_context

    def __isStopped(self):
        hc = self._jhc.h2oContext()
        field = hc.getClass().getSuperclass().getDeclaredField("stopped")
        field.setAccessible(True)
        return field.get(hc)

    def __isClientConnected(self):
        hc = self._jhc.h2oContext()
        field = self.__getClientConnectedField()
        return field.get(hc)

    def __setClientConnected(self):
        hc = self._jhc.h2oContext()
        field = self.__getClientConnectedField()
        field.set(hc, True)

    def __getClientConnectedField(self):
        hc = self._jhc.h2oContext()
        field = hc.getClass().getSuperclass().getDeclaredField("clientConnected")
        field.setAccessible(True)
        return field

    def stop(self):
        h2o.connection().close()
        hc = self._jhc.h2oContext()
        scalaStopMethod = getattr(hc, "org$apache$spark$h2o$H2OContext$$stop")
        scalaStopMethod(False, False, False) # stopSpark = False, stopJVM = False, inShutdownHook = False

        if self._conf.get("spark.ext.h2o.rest.api.based.client", "false") == "false":
            sys.exit()

    def download_h2o_logs(self, destination, container = "ZIP"):
        assert_is_type(container, Enum("ZIP", "LOG"))
        return self._jhc.h2oContext().downloadH2OLogs(destination, container)

    def __str__(self):
        if self.__isClientConnected() and not self.__isStopped():
            return self._jhc.toString()
        else:
            return "H2OContext has been stopped or hasn't been created. Call H2OContext.getOrCreate() or " \
                   "H2OContext.getOrCreate(conf) to create a new one."

    def __repr__(self):
        self.show()
        return ""

    def show(self):
        print(self)

    def get_conf(self):
        return self._conf

    def asSparkFrame(self, h2oFrame, copyMetadata=True):
        """
        Transforms given H2OFrame to Spark DataFrame

        Parameters
        ----------
         h2oFrame : H2OFrame
         copyMetadata: Bool = True

        Returns
        -------
         Spark DataFrame
        """

        if isinstance(h2oFrame, H2OFrame):
            frame_id = h2oFrame.frame_id
            jdf = self._jhc.asDataFrame(frame_id, copyMetadata)
            sqlContext = SparkSession.builder.getOrCreate()._wrapped
            df = DataFrame(jdf, sqlContext)
            # Attach h2o_frame to dataframe which forces python not to delete the frame when we leave the scope of this
            # method.
            # Without this, after leaving this method python would garbage collect the frame since it's not used
            # anywhere and spark. when executing any action on this dataframe, will fail since the frame
            # would be missing.
            df._h2o_frame = h2oFrame
            return df

    def as_spark_frame(self, h2o_frame, copy_metadata=True):
        warnings.warn("Method 'as_spark_frame' is deprecated and will be removed in release 3.30. Please use method 'asSparkFrame' instead!")
        return self.asSparkFrame(h2o_frame, copy_metadata)

    def asH2OFrame(self, sparkFrame, h2oFrameName=None, fullCols=-1):
        """
        Transforms given Spark RDD or DataFrame to H2OFrame.

        Parameters
        ----------
          sparkFrame : Spark RDD or DataFrame
          h2oFrameName : Optional name for resulting H2OFrame
          fullCols : number of first n columns which are sent to the client together with the data

        Returns
        -------
          H2OFrame which contains data of original input Spark data structure
        """
        if isinstance(sparkFrame, DataFrame):
            return fc._as_h2o_frame_from_dataframe(self, sparkFrame, h2oFrameName, fullCols)
        elif isinstance(sparkFrame, RDD) and sparkFrame.isEmpty():
            schema = StructType([])
            empty = self._spark_session.createDataFrame(self._spark_session.sparkContext.emptyRDD(), schema)
            return fc._as_h2o_frame_from_dataframe(self, empty, h2oFrameName, fullCols)
        elif isinstance(sparkFrame, RDD):
            # First check if the type T in RDD[T] is one of the python "primitive" types
            # String, Boolean, Int and Double (Python Long is converted to java.lang.BigInteger)
            if _is_of_simple_type(sparkFrame):
                first = _get_first(sparkFrame)
                # Make this code compatible with python 3.6 and python 2.7
                global long
                if sys.version_info > (3,):
                    long = int

                if isinstance(first, str):
                    return fc._as_h2o_frame_from_RDD_String(self, sparkFrame, h2oFrameName, fullCols)
                elif isinstance(first, bool):
                    return fc._as_h2o_frame_from_RDD_Bool(self, sparkFrame, h2oFrameName, fullCols)
                elif (isinstance(sparkFrame.min(), int) and isinstance(sparkFrame.max(), int)) or (isinstance(sparkFrame.min(), long) and isinstance(sparkFrame.max(), long)):
                    if sparkFrame.min() >= _jvm().Integer.MIN_VALUE and sparkFrame.max() <= _jvm().Integer.MAX_VALUE:
                        return fc._as_h2o_frame_from_RDD_Int(self, sparkFrame, h2oFrameName, fullCols)
                    elif sparkFrame.min() >= _jvm().Long.MIN_VALUE and sparkFrame.max() <= _jvm().Long.MAX_VALUE:
                        return fc._as_h2o_frame_from_RDD_Long(self, sparkFrame, h2oFrameName, fullCols)
                    else:
                        raise ValueError('Numbers in RDD Too Big')
                elif isinstance(first, float):
                    return fc._as_h2o_frame_from_RDD_Float(self, sparkFrame, h2oFrameName, fullCols)
            else:
                return fc._as_h2o_frame_from_complex_type(self, sparkFrame, h2oFrameName, fullCols)
        else:
            raise ValueError('The asH2OFrame method expects Spark DataFrame or RDD as the input only!')

    def as_h2o_frame(self, dataframe, framename=None, full_cols=-1):
        warnings.warn("Method 'as_h2o_frame' is deprecated and will be removed in release 3.30. Please use method 'asH2OFrame' instead!")
        return self.asH2OFrame(dataframe, framename, full_cols)
<|endoftext|>


<commit_before>from uber.common import *

__version__ = 'v0.1'

attendee_fields = [
    'full_name', 'first_name', 'last_name', 'email', 'zip_code', 'cellphone', 'ec_phone', 'badge_status_label', 'checked_in',
    'badge_type_label', 'ribbon_label', 'staffing', 'is_dept_head', 'assigned_depts_labels', 'weighted_hours', 'worked_hours',
    'badge_num'
]
fields = dict({
    'food_restrictions': {
        'sandwich_pref_labels': True,
        'standard_labels': True,
        'freeform': True
    },
    'shifts': {
        'worked_label': True,
        'job': ['type_label', 'location_label', 'name', 'description', 'start_time', 'end_time', 'extra15']
    }
}, **{field: True for field in attendee_fields})

class AttendeeLookup:
    def lookup(self, badge_num):
        with Session() as session:
            attendee = session.query(Attendee).filter_by(badge_num=badge_num).first()
            return attendee.to_dict(fields) if attendee else {'error': 'No attendee found with Badge #{}'.format(badge_num)}

    def search(self, query):
        with Session() as session:
            return [a.to_dict(fields) for a in session.search(query).all()]

services.register(AttendeeLookup(), 'attendee')

job_fields = dict({
    'name': True,
    'description': True,
    'location': True,
    'start_time': True,
    'end_time': True,
    'duration': True,
    'shifts': {
        'attendee': {
            'badge_num': True,
            'full_name': True,
            'first_name': True,
            'last_name': True,
            'email': True,
            'cellphone': True,
            'badge_printed_name': True
        }
    }

})

class JobLookup:
    def lookup(self, location):
        with Session() as session:
            columns = Job.__table__.columns
            location_column = columns['location']
            label_lookup = {val: key for key, val in location_column.type.choices.items()}
            return [job.to_dict(job_fields) for job in session.query(Job).filter_by(location=label_lookup[location]).all()]

services.register(JobLookup(), 'shifts')

class DepartmentLookup:
    def list(self):
        with Session() as session:
            output = {}
            for dept in c.JOB_LOCATION_VARS:
                output[dept] = dict(c.JOB_LOCATION_OPTS)[getattr(c, dept)]
            return output

services.register(DepartmentLookup(), 'dept')

config_fields = [
    'EVENT_NAME',
    'ORGANIZATION_NAME',
    'YEAR',
    'EPOCH',

    'EVENT_VENUE',
    'EVENT_VENUE_ADDRESS',

    'AT_THE_CON',
    'POST_CON',

]

class ConfigLookup:
    def info(self):
        output = {
            'API_VERSION': __version__
        }
        for field in event_fields:
            output[field] = getattr(c, field)
        return output

    def lookup(self, field):
        if field.upper() in event_fields:
            return getattr(c, field.upper())

services.register(ConfigLookup(), 'config')<commit_msg>PEP8!!!!!!!!!!!!<commit_after>from uber.common import *

__version__ = 'v0.1'

attendee_fields = [
    'full_name', 'first_name', 'last_name', 'email', 'zip_code', 'cellphone', 'ec_phone', 'badge_status_label', 'checked_in',
    'badge_type_label', 'ribbon_label', 'staffing', 'is_dept_head', 'assigned_depts_labels', 'weighted_hours', 'worked_hours',
    'badge_num'
]
fields = dict({
    'food_restrictions': {
        'sandwich_pref_labels': True,
        'standard_labels': True,
        'freeform': True
    },
    'shifts': {
        'worked_label': True,
        'job': ['type_label', 'location_label', 'name', 'description', 'start_time', 'end_time', 'extra15']
    }
}, **{field: True for field in attendee_fields})

class AttendeeLookup:
    def lookup(self, badge_num):
        with Session() as session:
            attendee = session.query(Attendee).filter_by(badge_num=badge_num).first()
            return attendee.to_dict(fields) if attendee else {'error': 'No attendee found with Badge #{}'.format(badge_num)}

    def search(self, query):
        with Session() as session:
            return [a.to_dict(fields) for a in session.search(query).all()]

services.register(AttendeeLookup(), 'attendee')

job_fields = dict({
    'name': True,
    'description': True,
    'location': True,
    'start_time': True,
    'end_time': True,
    'duration': True,
    'shifts': {
        'attendee': {
            'badge_num': True,
            'full_name': True,
            'first_name': True,
            'last_name': True,
            'email': True,
            'cellphone': True,
            'badge_printed_name': True
        }
    }

})

class JobLookup:
    def lookup(self, location):
        with Session() as session:
            columns = Job.__table__.columns
            location_column = columns['location']
            label_lookup = {val: key for key, val in location_column.type.choices.items()}
            return [job.to_dict(job_fields) for job in session.query(Job).filter_by(location=label_lookup[location]).all()]

services.register(JobLookup(), 'shifts')

class DepartmentLookup:
    def list(self):
        with Session() as session:
            output = {}
            for dept in c.JOB_LOCATION_VARS:
                output[dept] = dict(c.JOB_LOCATION_OPTS)[getattr(c, dept)]
            return output

services.register(DepartmentLookup(), 'dept')

config_fields = [
    'EVENT_NAME',
    'ORGANIZATION_NAME',
    'YEAR',
    'EPOCH',

    'EVENT_VENUE',
    'EVENT_VENUE_ADDRESS',

    'AT_THE_CON',
    'POST_CON',

]

class ConfigLookup:
    def info(self):
        output = {
            'API_VERSION': __version__
        }
        for field in event_fields:
            output[field] = getattr(c, field)
        return output

    def lookup(self, field):
        if field.upper() in event_fields:
            return getattr(c, field.upper())

services.register(ConfigLookup(), 'config')
<|endoftext|>