<commit_before>#!/usr/bin/env python

if __name__ == '__main__':
    import sys
    import os
    import gitautodeploy
    sys.stderr.write("\033[1;33m[WARNING]\033[0;33m GitAutoDeploy.py is deprecated. Please use \033[1;33m'python gitautodeploy%s'\033[0;33m instead.\033[0m\n" % (' ' + ' '.join(sys.argv[1:])).strip())
    gitautodeploy.main()
<commit_msg>Fix truncated command and options<commit_after>#!/usr/bin/env python

if __name__ == '__main__':
    import sys
    import os
    import gitautodeploy
    sys.stderr.write("\033[1;33m[WARNING]\033[0;33m GitAutoDeploy.py is deprecated. Please use \033[1;33m'python gitautodeploy%s'\033[0;33m instead.\033[0m\n" % (' ' + ' '.join(sys.argv[1:])).rstrip())
    gitautodeploy.main()
<|endoftext|>


<commit_before>import time
import re
from threading import Thread

# external py files
from modules.irc_socket import openSocket, sendMessage
from modules.irc_init import joinRoom
from bot import (get_user, get_message, local_time, basic_command, update_command,
	word_n, streamer_acorn, streamer_geek, roulette, check_int, get_int, uptime,
	followage, streamer, duel, quote, addquote, bttv_user_replace, raffle,
	check_int_in_string, give_points)
from modules.sql import (db_add_user,
	db_add_points_user,
	db_minus_points_user,
	db_add_points_global,
	db_check_user,
	db_get_points_user,
	db_get_points_user_first,
	db_get_points_user_int,
	db_get_user_rank)
from modules.temp import (cooldown, temp_opponent, duel_state, temp_user,
	raffle_state, raffle_entries, raffle_amount)
from modules.settings import twitch_irc
from modules.api import check_user_class, get_latest_follower, get_youtube_request

# connection to the irc server is created
s = openSocket()
joinRoom(s)
readbuffer = ""

while True:

		readbuffer = readbuffer + s.recv(1024).decode('UTF-8')
		temp = str.split(readbuffer, "\n")
		readbuffer = temp.pop()

		for line in temp:
			print(line)

			if "PING :tmi.twitch.tv" in line:
				response = "PONG :tmi.twitch.tv\r\n"
				print("[IRC] >>> ", response)
				bytes_response = str.encode(response)
				s.send(bytes_response)
				break
			user = get_user(line)
			message = get_message(line)
			print("[MESSAGE] >>> " + user + " typed :" + message)

			### TIMER ###

			class rouletteTimer(Thread):

				def run(self):
					#time.sleep(440)
					time.sleep(30)
					print("[INFO] >>> Removed from cooldown array user: {}".format(cooldown[0]))
					cooldown.pop(0)
					print("[INFO] >>> Users in cooldown ", cooldown)

			def run():
					print("[INFO] >>> Cooldown timer started for roulette user: {}".format(user))
					rouletteTimer().start()

			class raffleTimer(Thread):

				def run(self):
					print("raffle timer started")
					time.sleep(30)
					sendMessage(s, "The raffle ends in 30 seconds!")
					time.sleep(30)
					raffle_state = False
					print("[DEBUG] >>> raffle timer stopped")
					sendMessage(s, raffle())

			def raffle_run():
					raffleTimer().start()

			### STRING EXTRACTION ###
			try:
				char_1 = word_n(message, 0)
				char_2 = word_n(message, 1)
				char_3 = word_n(message, 2)
				points_duel = char_3
				quote_to_add = message[10:]
				only_int = get_int(message)
			except:
				pass

			#### POINT BASED COMMANDS ###

			# returns how many points a user owns
			if "!points" in message:
				db_add_user(user)
				sendMessage(s, db_get_points_user(user))

			if "!givepoints" in message:
				sendMessage(s, str(give_points(user, char_2, char_3)))

			if "!rank" in message:
				sendMessage(s, db_get_user_rank(user))
			if "!roulette" in message and user not in cooldown:
				if(check_int(char_2)):
					sendMessage(s, roulette(user, char_2))
					cooldown.append(user)
					print("[INFO] >>> ", cooldown)
					run()
				else:
					sendMessage(s, "You can only enter int values {} BabyRage".format(user))

			if "!raffle" in message:
				del raffle_entries[:]
				raffle_points = char_2
				if(check_user_class(user, "moderators")) and check_int(raffle_points):
					print("[DEBUG] >>> {} started a raffle for {} points".format(user, raffle_points))
					raffle_state = True
					raffle_amount.append(raffle_points)
					sendMessage(s, "A raffle has started for {} points! Type !join to enter PogChamp".format(raffle_points))
					raffle_run()

			if "!join" in message and raffle_state:
				raffle_entries.append(user)
				print("[DEBUG >>> Users in raffle entries >>> ", raffle_entries)
			if "!join" in message and raffle_state == False:
				sendMessage(s, "{}, there is currently no active raffle BabyRage".format(user))

			if "!test" in message:
				sendMessage(s, "test message recieved")

			if "!duel" in message:
				test = re.match('(\w+\s\w+)', message[6:])
				if(test) == None:
					sendMessage(s, "You didn't specify an amount {}! FailFish".format(user))
					break
				if(check_int(points_duel)):

					check_user_points = db_get_points_user_int(user)
					if(check_user_points < int(points_duel)):
						sendMessage(s, "Sorry {}, you don't have enough points for that BabyRage".format(user))
						break

					opponent = char_2
					format_opponent = bttv_user_replace(opponent)
					temp_opponent.append(opponent.lower())
					temp_user.append(user.lower())
					duel_state = True

					if(format_opponent == twitch_irc.get('NICK')):
						sendMessage(s, "{}, I always win the duel! MingLee".format(user))
						break

					sendMessage(s, "{}, you have been challenged to {} points by {}! Type !accept to duel.".format(opponent, points_duel, user))
				else:
					sendMessage(s, "{} you can only enter integers! BabyRage".format(user))

					print(points_duel)
					print(temp_opponent)
					print(duel_state)

			if "!accept" in message and duel_state and (user in temp_user or user in temp_opponent):

				print("[DEBUG] >>> Points to duel >>> ", points_duel)
				print("[DEBUG] >>> Users in temp_opponent >>> {}".format(temp_opponent))
				try:
					# !duel pokchok 10
					user_to_duel = temp_opponent[0] # pokchok
					original_duel_user = temp_user[0] # kritzware
				except Exception as e:
					print(e)
					sendMessage(s, "{} you don't currently have an active duel Kappa".format(user))
					break

				print(user_to_duel)
				print(original_duel_user)

				if(user in temp_opponent):
					print("[DEBUG] >>> Checking {} in temp_opponents".format(user))
					print(original_duel_user + " duelling " + user_to_duel)
					sendMessage(s, duel(original_duel_user, user_to_duel, points_duel))
					temp_opponent.remove(user_to_duel)
					temp_user.remove(original_duel_user)
					duel_state = False
					points_duel = ''
				else:
					print("[ERROR] >>> {} not found in temp_opponent".format(user))

			# if "!songrequest" in message:
			# 	sendMessage(s, get_youtube_request(user, char_2))

			### ADVANCED COMMANDS ###
			if "!uptime" in message:
				print("[COMMAND] >>> !uptime")
				sendMessage(s, uptime())
			if "!followage" in message:
				print("[COMMAND] >>> !followage")
				sendMessage(s, followage(user))
			if "!top" in message:
				print("[COMMAND] >>> !top")
				sendMessage(s, db_get_points_user_first())
			if "!quote" in message:
				print("[COMMAND] >>> !quote")
				sendMessage(s, quote())
			if "!addquote" in message:
				print("[COMMAND] >>> !addquote")
				sendMessage(s, addquote(quote_to_add))

			if "!followage" in message:
				print("[COMMAND] >>> !followage")
				sendMessage(s, followage(user))

			if "!streamer" in message:
				sendMessage(s, streamer(user, char_2))

			### DEFAULT COMMANDS ###
			if "!localtime" in message:
				print("[COMMAND] >>> !localtime")
				sendMessage(s, local_time())
			if "!admin" in message:
				print("[COMMAND] >>> !admin")
				sendMessage(s, basic_command('admin', user))
			if "!help" in message or "!commands" in message:
				print("[COMMAND] >>> !help")
				sendMessage(s, basic_command('help', user))
			if "!twitter" in message:
				print("[COMMAND] >>> !twitter")
				sendMessage(s, basic_command('twitter', user))
			if "!spooky" in message:
				print("[COMMAND] >>> !spooky")
				sendMessage(s, basic_command('spooky', user))
			if "!bot" in message:
				print("[COMMAND] >>> !bot")
				sendMessage(s, basic_command('bot', user))
			if "!donger" in message:
				print("[COMMAND] >>> !donger")
				try:
					sendMessage(s, basic_command('donger', user))
				except:
					pass

			# ### STREAMERS###
			# if "!streamer acorn" in message:
			# 	sendMessage(s, streamer_acorn(user))
			# if "!streamer geek" in message:
			# 	sendMessage(s, streamer_geek(user))

			### UPDATE COMMANDS ###
			if "!coloring" in message:
				print("[COMMAND] >>> !coloring")
				sendMessage(s, update_command('coloring', user))

			if "!splinter" in message:
				sendMessage(s, "NO FUN ALLOWED FUNgineer NO SONG REQUESTS FUNgineer")<commit_msg>Added timer to remove users from a duel if no response after 60 secs<commit_after>import time
import re
from threading import Thread

# external py files
from modules.irc_socket import openSocket, sendMessage
from modules.irc_init import joinRoom
from bot import (get_user, get_message, local_time, basic_command, update_command,
	word_n, streamer_acorn, streamer_geek, roulette, check_int, get_int, uptime,
	followage, streamer, duel, quote, addquote, bttv_user_replace, raffle,
	check_int_in_string, give_points)
from modules.sql import (db_add_user,
	db_add_points_user,
	db_minus_points_user,
	db_add_points_global,
	db_check_user,
	db_get_points_user,
	db_get_points_user_first,
	db_get_points_user_int,
	db_get_user_rank)
from modules.temp import (cooldown, temp_opponent, duel_state, temp_user,
	raffle_state, raffle_entries, raffle_amount)
from modules.settings import twitch_irc
from modules.api import check_user_class, get_latest_follower, get_youtube_request

# connection to the irc server is created
s = openSocket()
joinRoom(s)
readbuffer = ""

while True:

		readbuffer = readbuffer + s.recv(1024).decode('UTF-8')
		temp = str.split(readbuffer, "\n")
		readbuffer = temp.pop()

		for line in temp:
			print(line)

			if "PING :tmi.twitch.tv" in line:
				response = "PONG :tmi.twitch.tv\r\n"
				print("[IRC] >>> ", response)
				bytes_response = str.encode(response)
				s.send(bytes_response)
				break
			user = get_user(line)
			message = get_message(line)
			print("[MESSAGE] >>> " + user + " typed :" + message)

			### TIMER ###

			class rouletteTimer(Thread):

				def run(self):
					#time.sleep(440)
					time.sleep(30)
					print("[INFO] >>> Removed from cooldown array user: {}".format(cooldown[0]))
					cooldown.pop(0)
					print("[INFO] >>> Users in cooldown ", cooldown)

			def run():
					print("[INFO] >>> Cooldown timer started for roulette user: {}".format(user))
					rouletteTimer().start()

			class raffleTimer(Thread):

				def run(self):
					print("[INFO] >>> raffle timer started")
					time.sleep(30)
					sendMessage(s, "The raffle ends in 30 seconds!")
					time.sleep(30)
					raffle_state = False
					print("[DEBUG] >>> raffle timer stopped")
					sendMessage(s, raffle())

			def raffle_run():
					raffleTimer().start()

			class duelTimer(Thread):
				def run(self):
					print("[INFO] >>> duel timer started")
					print(temp_opponent)
					print(temp_user)
					time.sleep(60)
					print("[INFO] >>> {} removed from duel queue")
					temp_opponent.pop(0)

			def duel_run():
					duelTimer().start()

			### STRING EXTRACTION ###
			try:
				char_1 = word_n(message, 0)
				char_2 = word_n(message, 1)
				char_3 = word_n(message, 2)
				points_duel = char_3
				quote_to_add = message[10:]
				only_int = get_int(message)
			except:
				pass

			#### POINT BASED COMMANDS ###

			# returns how many points a user owns
			if "!points" in message:
				db_add_user(user)
				sendMessage(s, db_get_points_user(user))

			if "!givepoints" in message:
				sendMessage(s, str(give_points(user, char_2, char_3)))

			if "!rank" in message:
				sendMessage(s, db_get_user_rank(user))
			if "!roulette" in message and user not in cooldown:
				if(check_int(char_2)):
					sendMessage(s, roulette(user, char_2))
					cooldown.append(user)
					print("[INFO] >>> ", cooldown)
					run()
				else:
					sendMessage(s, "You can only enter int values {} BabyRage".format(user))

			if "!raffle" in message:
				#if(raffle_state == True):
				#	sendMessage(s, "A raffle is already active OpieOP")
				if(raffle_state == False):
					del raffle_entries[:]
					raffle_points = char_2
					if(check_user_class(user, "moderators")) and check_int(raffle_points):
						print("[DEBUG] >>> {} started a raffle for {} points".format(user, raffle_points))
						raffle_state = True
						raffle_amount.append(raffle_points)
						sendMessage(s, "A raffle has started for {} points! Type !join to enter PogChamp".format(raffle_points))
						raffle_run()

			if "!join" in message and raffle_state:
				raffle_entries.append(user)
				print("[DEBUG >>> Users in raffle entries >>> ", raffle_entries)
			if "!join" in message and raffle_state == False:
				sendMessage(s, "{}, there is currently no active raffle BabyRage".format(user))

			if "!test" in message:
				print(temp_opponent)
				print(temp_user)

			if "!duel" in message:
				test = re.match('(\w+\s\w+)', message[6:])
				if(test) == None:
					sendMessage(s, "You didn't specify an amount {}! FailFish".format(user))
					break
				if(check_int(points_duel)):

					check_user_points = db_get_points_user_int(user)
					if(check_user_points < int(points_duel)):
						sendMessage(s, "Sorry {}, you don't have enough points for that BabyRage".format(user))
						break

					opponent = char_2
					format_opponent = bttv_user_replace(opponent)
					temp_opponent.append(opponent.lower())
					temp_user.append(user.lower())
					duel_state = True

					if(format_opponent == twitch_irc.get('NICK')):
						sendMessage(s, "{}, I always win the duel! MingLee".format(user))
						break

					duel_run()
					sendMessage(s, "{}, you have been challenged to {} points by {}! Type !accept to duel.".format(opponent, points_duel, user))
				else:
					sendMessage(s, "{} you can only enter integers! BabyRage".format(user))

					print(points_duel)
					print(temp_opponent)
					print(duel_state)

			if "!accept" in message and duel_state and (user in temp_user or user in temp_opponent):

				print("[DEBUG] >>> Points to duel >>> ", points_duel)
				print("[DEBUG] >>> Users in temp_opponent >>> {}".format(temp_opponent))
				try:
					# !duel pokchok 10
					user_to_duel = temp_opponent[0] # pokchok
					original_duel_user = temp_user[0] # kritzware
				except Exception as e:
					print(e)
					sendMessage(s, "{} you don't currently have an active duel Kappa".format(user))
					break

				print(user_to_duel)
				print(original_duel_user)

				if(user in temp_opponent):
					print("[DEBUG] >>> Checking {} in temp_opponents".format(user))
					print(original_duel_user + " duelling " + user_to_duel)
					sendMessage(s, duel(original_duel_user, user_to_duel, points_duel))
					temp_opponent.remove(user_to_duel)
					temp_user.remove(original_duel_user)
					duel_state = False
					points_duel = ''
				else:
					print("[ERROR] >>> {} not found in temp_opponent".format(user))

			# if "!songrequest" in message:
			# 	sendMessage(s, get_youtube_request(user, char_2))

			### ADVANCED COMMANDS ###
			if "!uptime" in message:
				print("[COMMAND] >>> !uptime")
				sendMessage(s, uptime())
			if "!followage" in message:
				print("[COMMAND] >>> !followage")
				sendMessage(s, followage(user))
			if "!top" in message:
				print("[COMMAND] >>> !top")
				sendMessage(s, db_get_points_user_first())
			if "!quote" in message:
				print("[COMMAND] >>> !quote")
				sendMessage(s, quote())
			if "!addquote" in message:
				print("[COMMAND] >>> !addquote")
				sendMessage(s, addquote(quote_to_add))

			if "!followage" in message:
				print("[COMMAND] >>> !followage")
				sendMessage(s, followage(user))

			if "!streamer" in message:
				sendMessage(s, streamer(user, char_2))

			### DEFAULT COMMANDS ###
			if "!localtime" in message:
				print("[COMMAND] >>> !localtime")
				sendMessage(s, local_time())
			if "!admin" in message:
				print("[COMMAND] >>> !admin")
				sendMessage(s, basic_command('admin', user))
			if "!help" in message or "!commands" in message:
				print("[COMMAND] >>> !help")
				sendMessage(s, basic_command('help', user))
			if "!twitter" in message:
				print("[COMMAND] >>> !twitter")
				sendMessage(s, basic_command('twitter', user))
			if "!spooky" in message:
				print("[COMMAND] >>> !spooky")
				sendMessage(s, basic_command('spooky', user))
			if "!bot" in message:
				print("[COMMAND] >>> !bot")
				sendMessage(s, basic_command('bot', user))
			if "!donger" in message:
				print("[COMMAND] >>> !donger")
				try:
					sendMessage(s, basic_command('donger', user))
				except:
					pass

			# ### STREAMERS###
			# if "!streamer acorn" in message:
			# 	sendMessage(s, streamer_acorn(user))
			# if "!streamer geek" in message:
			# 	sendMessage(s, streamer_geek(user))

			### UPDATE COMMANDS ###
			if "!coloring" in message:
				print("[COMMAND] >>> !coloring")
				sendMessage(s, update_command('coloring', user))

			if "!splinter" in message:
				sendMessage(s, "NO FUN ALLOWED FUNgineer NO SONG REQUESTS FUNgineer")<|endoftext|>


<commit_before>#!/usr/bin/env python
# Program Name : js_template.py
# Description  : Delphix implementation script
# Author       : Corey Brune
# Created: March 4 2016
#
# Copyright (c) 2016 by Delphix.
# All rights reserved.
# See http://docs.delphix.com/display/PS/Copyright+Statement for details
#
# Delphix Support statement available at
# See http://docs.delphix.com/display/PS/PS+Script+Support+Policy for details
#
# Warranty details provided in external file
# for customers who have purchased support.
#
"""Refresh a vdb
Usage:
  js_template.py (--create_template <template_name> --database <name> | --list_templates | --delete_template <template_name>)
                   [--engine <identifier> | --all] [--parallel <n>]
                   [--poll <n>] [--debug]
                   [--config <path_to_file>] [--logdir <path_to_file>]
  js_template.py -h | --help | -v | --version

Creates, Lists, Removes a Jet Stream Template

Examples:
  js_template.py --list-template
  js_template.py --create-template <template_name> --database <name>
  js_template.py --create-template <template_name> --database <name:name:name>
  js_template.py --delete-template <template_name>

Options:
  --create-template <name>  Name of the new JS Template
  --database <name>         Name of the database(s) to use for the JS Template
                                Note: If adding multiple template DBs, use a
                                comma (:) to delineate between the DB names.
  --list_templates          List the templates on a given engine
  --engine <type>           Alt Identifier of Delphix engine in dxtools.conf.
  --all                     Run against all engines.
  --debug                   Enable debug logging
  --parallel <n>            Limit number of jobs to maxjob
  --poll <n>                The number of seconds to wait between job polls
                            [default: 10]
  --config <path_to_file>   The path to the dxtools.conf file
                            [default: ./dxtools.conf]
  --logdir <path_to_file>   The path to the logfile you want to use.
                            [default: ./js_template.log]
  -h --help                 Show this screen.
  -v --version              Show version.
"""

VERSION="v.0.0.001"

from docopt import docopt
import logging
from os.path import basename
import signal
import sys
import time
import traceback
import json
import re
from datetime import datetime
from dateutil import tz

from multiprocessing import Process
from time import sleep, time
from delphixpy.v1_6_0.delphix_engine import DelphixEngine
from delphixpy.v1_6_0.web.jetstream import template
from delphixpy.v1_6_0 import job_context
from delphixpy.v1_6_0.web import database
from delphixpy.v1_6_0.web.service import time
from delphixpy.v1_6_0.web.vo import JSDataTemplateCreateParameters
from delphixpy.v1_6_0.exceptions import RequestError
from delphixpy.v1_6_0.exceptions import JobError
from delphixpy.v1_6_0.exceptions import HttpError

from lib.DlpxException import DlpxException
from lib.DxTimeflow import DxTimeflow
from lib.GetReferences import convert_timestamp
from lib.GetReferences import find_obj_by_name
from lib.DxLogging import print_error
from lib.GetSession import GetSession

def create_template(template_name, database_name):
    assert template_name and database_name is not '', 'The template ' \
            'and database names are required for creating a JS Template.\n'

    js_template_params = JSDataTemplateCreateParameters()

    template_ds_lst = []
    for db in database_name.split(','):
            template_ds_lst.append(build_ds_params(engine, db))

    try:
        js_template_params.data_sources = template_ds_lst
        js_template_params.name = template_name
        js_template_params.type = 'JSDataTemplateCreateParameters'

        template_ret_val = template.create(dx_session_obj.server_session,
                                           js_template_params)

        print('Template %s was created successfully with reference %s\n' %
              (template_name, template_ret_val))

    except DlpxException, e:
        print('\nThe template %s was not created. The error was:\n\n%s' %
              (template_name, e.errors))
        sys.exit(1)

    except AssertionError, e:
        print('\nAn error occurred creating the container, %s:\n\n%s' %
              (container_name, e))

def list_template(engine):
    """
    List all templates on a given engine

    engine: Delphix Engine session object
    """

    js_list = ['name','reference','active_branch','last_updated']
    header = 'Name\t\t\tReference\t\tActive Branch\tLast Updated'
    print header

    try:
        js_templates = template.get_all(engine)

        for js_template in js_templates:
            last_updated = convertTimestamp(engine,
                                            js_template.last_updated[:-5])

        print js_template.name, js_template.reference, \
            js_template.active_branch, last_updated

    except (DlpxException, HttpError, RequestError) as e:
        raise DlpxException('\nERROR: The templates on engine %s could not '
                            'be listed. The error was:\n\n%s' %
                            (engine.ip_addr, e.message))

def delete_template(engine, template_name):
    """
    Deletes a template

    engine: Delphix Engine session object
    template_name: Template to delete
    """

    try:
        template_obj = find_obj_by_name(engine, template, template_name)
        template.delete(engine, template_obj.reference)

        print 'Template %s is deleted.' % (template_name)

    except (DlpxException, HttpError, RequestError) as e:
        raise DlpxException('\nERROR: The template %s was not deleted. The '
                            'error was:\n\n%s' % (template_name, e.message))

def build_ds_params(engine, db):
    try:
        db_obj = find_obj_by_name(engine, db)
        return({'type': 'JSDataSourceCreateParameters', 'source': {'type':
                'JSDataSource', 'name': db}, 'container': db_obj.reference})

    except RequestError as e:
        print('\nCould not find %s\n%s' % (db, e.message))
        sys.exit(1)

def job_mode(server):
    """
    This function tells Delphix how to execute jobs, based on the
    single_thread variable at the beginning of the file
    """
    #Synchronously (one at a time)
    if single_thread == True:
        job_m = job_context.sync(server)
        print_debug("These jobs will be executed synchronously")
    #Or asynchronously
    else:
        job_m = job_context.async(server)
        print_debug("These jobs will be executed asynchronously")
    return job_m

@run_async
def main_workflow(engine):
    """
    This function is where we create our main workflow.
    Use the @run_async decorator to run this function asynchronously.
    The @run_async decorator allows us to run against multiple Delphix Engine
    simultaneously

    engine: Dictionary containing engine IP, username, password
    """

    #Establish these variables as empty for use later
    environment_obj = None
    source_objs = None
    jobs = {}

    try:
        #Setup the connection to the Delphix Engine
        dx_session_obj.serversess(engine['ip_address'], engine['username'],
                                  engine['password'])

    except DlpxException as e:
        raise DlpxException(e)

    thingstodo = ["thingtodo"]
    #reset the running job count before we begin
    i = 0

    try:
        with job_mode(server):
            while (len(jobs) > 0 or len(thingstodo)> 0):
                if len(thingstodo)> 0:

                    if arguments['--create_template']:
                        create_template(arguments['--create_template'],
                                        arguments['--database'])
                    elif arguments['--delete_template']:
                    elif arguments['--list_template']:

                thingstodo.pop()
                #get all the jobs, then inspect them
                i = 0
                for j in jobs.keys():
                    job_obj = job.get(server, jobs[j])
                    print_debug(job_obj)
                    if job_obj.job_state in ["CANCELED", "COMPLETED", "FAILED"]:
                        #If the job is in a non-running state, remove it
                        #from the running jobs list.
                        del jobs[j]
                    else:
                        #If the job is in a running state, increment the
                        #running job count.
                        i += 1

            print_info('\n%s: %s jobs running.\n' % (engine['hostname'],
                       str(i)))

            #If we have running jobs, pause before repeating the checks.
            if len(jobs) > 0:
                sleep(float(arguments['--poll']))

    except (DlpxException, JobError) as e:
        print('\nERROR: %s engine encountered an error:\n%s' %
              (database_name, e.message))
        sys.exit(1)

def run_job():
    """
    This function runs the main_workflow aynchronously against all the
    servers specified

    No arguments. run_job() uses global variables instead of args.
    """

    #Create an empty list to store threads we create.
    threads = []

    #Else if the --engine argument was given, test to see if the engine
    # exists in dxtools.conf
    if arguments['--engine']:
        try:
            engine = dx_session_obj.dlpx_engines[arguments['--engine']]
            print_info("Executing against Delphix Engine: " +
                       arguments['--engine'])

        except KeyError:
            raise DlpxException('\nERROR: Delphix Engine %s cannot be '
                                'found in %s.\nPlease check your values '
                                'and try again. Exiting\n' %
                                (arguments['--engine'],
                                config_file_path))

    elif arguments['--engine'] is None:
        #search for a default engine in the dxtools.conf
        for delphix_engine in dx_session_obj.dlpx_engines:

            if dx_session_obj.dlpx_engines[delphix_engine]['default'] == 'true':
                engine = dx_session_obj.dlpx_engines[delphix_engine]
                print_info('Executing against the default Delphix Engine in '
                       'the dxtools.conf: %s' % (
                       dx_session_obj.dlpx_engines[delphix_engine]['hostname']))

                break

        if engine is None:
            raise DlpxException("\nNo engine specified with the --engine arg '
                                ' and no default engine found. Exiting.\n")

        #run the job against the engine
        main_workflow(engine)

def main(argv):
    #We want to be able to call on these variables anywhere in the script.
    global single_thread
    global usebackup
    global time_start
    global host_name
    global database_name
    global config_file_path
    global dx_session_obj

    try:
        dx_session_obj = GetSession()
        logging_est(arguments['--logdir'])
        print_debug(arguments)
        time_start = time()
        database_name = arguments['--vdb']
        config_file_path = arguments['--config']

        logging_est(arguments['--logdir'])
        print_debug(arguments)
        time_start = time()
        engine = None
        single_thread = False
        database_name = arguments['--name']
        host_name = arguments['--host']
        config_file_path = arguments['--config']
        #Parse the dxtools.conf and put it into a dictionary
        dxtools_objects = get_config(config_file_path)

        #This is the function that will handle processing main_workflow for
        # all the servers.
        run_job(engine)

        elapsed_minutes = time_elapsed()
        print_info("script took " + str(elapsed_minutes) +
                   " minutes to get this far.")

    #Here we handle what we do when the unexpected happens
    except SystemExit as e:
        """
        This is what we use to handle our sys.exit(#)
        """
        sys.exit(e)
    except HttpError as e:
        """
        We use this exception handler when our connection to Delphix fails
        """
        print_error("Connection failed to the Delphix Engine")
        print_error( "Please check the ERROR message below")
        print_error(e.message)
        sys.exit(2)
    except JobError as e:
        """
        We use this exception handler when a job fails in Delphix so that we
        have actionable data
        """
        print_error("A job failed in the Delphix Engine")
        print_error(e.job)
        elapsed_minutes = time_elapsed()
        print_info(basename(__file__) + " took " + str(elapsed_minutes) +
                   " minutes to get this far.")
        sys.exit(3)
    except KeyboardInterrupt:
        """
        We use this exception handler to gracefully handle ctrl+c exits
        """
        print_debug("You sent a CTRL+C to interrupt the process")
        elapsed_minutes = time_elapsed()
        print_info(basename(__file__) + " took " + str(elapsed_minutes) +
                   " minutes to get this far.")
    except:
        """
        Everything else gets caught here
        """
        print_error(sys.exc_info()[0])
        print_error(traceback.format_exc())
        elapsed_minutes = time_elapsed()
        print_info(basename(__file__) + " took " + str(elapsed_minutes) +
                   " minutes to get this far.")
        sys.exit(1)

if __name__ == "__main__":
    #Grab our arguments from the doc at the top of the script
    arguments = docopt(__doc__, version=basename(__file__) + " " + VERSION)

    #Feed our arguments to the main function, and off we go!
    main(arguments)
<commit_msg>Delete js_template.py<commit_after><|endoftext|>


<commit_before>#!/usr/bin/env python

from Tkinter import *
from requests import Session
import json
import sys
from PIL import Image, ImageTk
from StringIO import StringIO
import os

class HSDeckBuilder(Frame):
    def __init__(self, session, master=None):
        Frame.__init__(self, master)
        self.master.title('HSDeckBuilder')
        self.create_widgets(session)

    def create_widgets(self, session):
        self.card_spot_list = []
        for x in range(0, 4):
            card = Label(text="Card_Spot", image=None)
            card.grid(column=x, row=1)
            self.card_spot_list.append(card)
        for x in range(0, 4):
            card = Label(text="Card_Spot", image=None)
            card.grid(column=x, row=2)
            self.card_spot_list.append(card)

        self.druid_button = Button(text="Druid", image=None)
        self.druid_button.grid(column=0, row=0)
        self.druid_button.bind('<Button-1>',
                lambda event: load_cards(session, 'Mage', self.card_spot_list))

def load_cards(session, playerClass, card_spot_list):
    card_data = open("AllSets.enUS.json").read()
    card_data = json.loads(card_data)
    count = 0
    for item in card_data['Expert']:
        if 'playerClass' in item:
            if item['playerClass'] == playerClass:
                print item['name'], ' ', item['id']
                url = r'http://wow.zamimg.com/images/hearthstone/cards/enus/original/%s.png' % item['id']
                display_image(session, url, card_spot_list, count)
                count = count + 1
                if count == len(card_spot_list):
                    return

def display_image(session, url, card_spot_list, count):
    print 'click'
    response = session.get(url)
    img = Image.open(StringIO(response.content))
    img2 = ImageTk.PhotoImage(img)
    card_spot_list[count].config(image = img2)
    card_spot_list[count].image = img2

#Main application entry point
def main():
    session = Session()
    session.headers = {'user-agent': 'HearthStone Deck Builder v0.1'}
    app = HSDeckBuilder(session)
    app.mainloop()
#Standard main function
if __name__ == '__main__':
       main()<commit_msg>Click buttons to display first 8 alphabetical cards from different classes. A lot of valuesa re currently hardcoded in<commit_after>#!/usr/bin/env python

from Tkinter import *
from requests import Session
import json
import sys
from PIL import Image, ImageTk
from StringIO import StringIO
import os

HERO_CLASSES = ['Druid', 'Mage', 'Warlock', 'Hunter', 'Priest', 'Paladin', 'Rogue']

class HSDeckBuilder(Frame):
    def __init__(self, session, master=None):
        Frame.__init__(self, master)
        self.master.title('HSDeckBuilder')
        self.create_widgets(session)

    def create_widgets(self, session):
        self.card_spot_list = []
        for x in range(0, 4):
            card = Label(text="Card_Spot", image=None)
            card.grid(column=x, row=1)
            self.card_spot_list.append(card)
        for x in range(0, 4):
            card = Label(text="Card_Spot", image=None)
            card.grid(column=x, row=2)
            self.card_spot_list.append(card)

        for x, name in enumerate(HERO_CLASSES):
            button = Button(text=name, image=None)
            button.grid(column=x, row=0)
            def handler(event, session=session, name=name, card_spot_list=self.card_spot_list):
                return load_cards(event, session, name, self.card_spot_list)
            button.bind('<Button-1>', handler)

def load_cards(event, session, playerClass, card_spot_list):
    card_data = open("AllSets.enUS.json").read()
    card_data = json.loads(card_data)
    count = 0
    for item in card_data['Expert']:
        if 'playerClass' in item:
            if item['playerClass'] == playerClass:
                print item['name'], ' ', item['id']
                url = r'http://wow.zamimg.com/images/hearthstone/cards/enus/original/%s.png' % item['id']
                response = session.get(url)
                if response.status_code != 200:
                    continue
                display_image(session, url, card_spot_list, count)
                count = count + 1
                if count == len(card_spot_list):
                    return

def display_image(session, url, card_spot_list, count):
    response = session.get(url)
    img = Image.open(StringIO(response.content)).resize((200, 304),Image.ANTIALIAS)
    img2 = ImageTk.PhotoImage(img)
    card_spot_list[count].config(image = img2)
    card_spot_list[count].image = img2

#Main application entry point
def main():
    print HERO_CLASSES
    session = Session()
    session.headers = {'user-agent': 'HearthStone Deck Builder v0.1'}
    app = HSDeckBuilder(session)
    app.mainloop()
#Standard main function
if __name__ == '__main__':
       main()<|endoftext|>


<commit_before>"""
Test the pipeline module.
"""
import numpy as np
from scipy import sparse

from nose.tools import assert_raises, assert_equal, assert_false, assert_true
from numpy.testing import assert_array_equal, \
        assert_array_almost_equal

from sklearn.base import BaseEstimator, clone
from sklearn.pipeline import Pipeline, FeatureStacker
from sklearn.svm import SVC
from sklearn.linear_model import LogisticRegression
from sklearn.feature_selection import SelectKBest, f_classif
from sklearn.decomposition.pca import PCA, RandomizedPCA
from sklearn.datasets import load_iris
from sklearn.preprocessing import StandardScaler
from sklearn.feature_extraction.text import CountVectorizer

class IncorrectT(BaseEstimator):
    """Small class to test parameter dispatching.
    """

    def __init__(self, a=None, b=None):
        self.a = a
        self.b = b

class T(IncorrectT):

    def fit(self, X, y):
        return self

class TransfT(T):

    def transform(self, X, y=None):
        return X

class FitParamT(BaseEstimator):
    """Mock classifier
    """

    def __init__(self):
        self.successful = False
        pass

    def fit(self, X, y, should_succeed=False):
        self.successful = should_succeed

    def predict(self, X):
        return self.successful

def test_pipeline_init():
    """ Test the various init parameters of the pipeline.
    """
    assert_raises(TypeError, Pipeline)
    # Check that we can't instantiate pipelines with objects without fit
    # method
    pipe = assert_raises(TypeError, Pipeline,
                        [('svc', IncorrectT)])
    # Smoke test with only an estimator
    clf = T()
    pipe = Pipeline([('svc', clf)])
    assert_equal(pipe.get_params(deep=True),
                 dict(svc__a=None, svc__b=None, svc=clf))

    # Check that params are set
    pipe.set_params(svc__a=0.1)
    assert_equal(clf.a, 0.1)
    # Smoke test the repr:
    repr(pipe)

    # Test with two objects
    clf = SVC()
    filter1 = SelectKBest(f_classif)
    pipe = Pipeline([('anova', filter1), ('svc', clf)])

    # Check that params are set
    pipe.set_params(svc__C=0.1)
    assert_equal(clf.C, 0.1)
    # Smoke test the repr:
    repr(pipe)

    # Check that params are not set when naming them wrong
    assert_raises(ValueError, pipe.set_params, anova__C=0.1)

    # Test clone
    pipe2 = clone(pipe)
    assert_false(pipe.named_steps['svc'] is pipe2.named_steps['svc'])

    # Check that appart from estimators, the parameters are the same
    params = pipe.get_params()
    params2 = pipe2.get_params()
    # Remove estimators that where copied
    params.pop('svc')
    params.pop('anova')
    params2.pop('svc')
    params2.pop('anova')
    assert_equal(params, params2)

def test_pipeline_methods_anova():
    """ Test the various methods of the pipeline (anova).
    """
    iris = load_iris()
    X = iris.data
    y = iris.target
    # Test with Anova + LogisticRegression
    clf = LogisticRegression()
    filter1 = SelectKBest(f_classif, k=2)
    pipe = Pipeline([('anova', filter1), ('logistic', clf)])
    pipe.fit(X, y)
    pipe.predict(X)
    pipe.predict_proba(X)
    pipe.predict_log_proba(X)
    pipe.score(X, y)

def test_pipeline_fit_params():
    """Test that the pipeline can take fit parameters
    """
    pipe = Pipeline([('transf', TransfT()), ('clf', FitParamT())])
    pipe.fit(X=None, y=None, clf__should_succeed=True)
    # classifier should return True
    assert_true(pipe.predict(None))
    # and transformer params should not be changed
    assert_true(pipe.named_steps['transf'].a is None)
    assert_true(pipe.named_steps['transf'].b is None)

def test_pipeline_methods_pca_svm():
    """Test the various methods of the pipeline (pca + svm)."""
    iris = load_iris()
    X = iris.data
    y = iris.target
    # Test with PCA + SVC
    clf = SVC(probability=True)
    pca = PCA(n_components='mle', whiten=True)
    pipe = Pipeline([('pca', pca), ('svc', clf)])
    pipe.fit(X, y)
    pipe.predict(X)
    pipe.predict_proba(X)
    pipe.predict_log_proba(X)
    pipe.score(X, y)

def test_pipeline_methods_preprocessing_svm():
    """Test the various methods of the pipeline (preprocessing + svm)."""
    iris = load_iris()
    X = iris.data
    y = iris.target
    n_samples = X.shape[0]
    n_classes = len(np.unique(y))
    scaler = StandardScaler()
    pca = RandomizedPCA(n_components=2, whiten=True)
    clf = SVC(probability=True)

    for preprocessing in [scaler, pca]:
        pipe = Pipeline([('scaler', scaler), ('svc', clf)])
        pipe.fit(X, y)

        # check shapes of various prediction functions
        predict = pipe.predict(X)
        assert_equal(predict.shape, (n_samples,))

        proba = pipe.predict_proba(X)
        assert_equal(proba.shape, (n_samples, n_classes))

        log_proba = pipe.predict_log_proba(X)
        assert_equal(log_proba.shape, (n_samples, n_classes))

        decision_function = pipe.decision_function(X)
        assert_equal(decision_function.shape, (n_samples, n_classes))

        pipe.score(X, y)

def test_feature_stacker():
    # basic sanity check for feature stacker
    iris = load_iris()
    X = iris.data
    X -= X.mean(axis=0)
    y = iris.target
    pca = RandomizedPCA(n_components=2)
    select = SelectKBest(k=1)
    fs = FeatureStacker([("pca", pca), ("select", select)])
    fs.fit(X, y)
    X_transformed = fs.transform(X)
    assert_equal(X_transformed.shape, (X.shape[0], 3))

    # check if it does the expected thing
    assert_array_almost_equal(X_transformed[:, :-1], pca.fit_transform(X))
    assert_array_equal(X_transformed[:, -1],
            select.fit_transform(X, y).ravel())

    # test if it also works for sparse input
    X_sp = sparse.csr_matrix(X)
    X_sp_transformed = fs.fit_transform(X_sp, y)
    assert_array_almost_equal(X_transformed, X_sp_transformed.toarray())

    # test setting parameters
    fs.set_params(select__k=2)
    assert_equal(fs.fit_transform(X, y).shape, (X.shape[0], 4))

def test_feature_stacker_feature_names():
    JUNK_FOOD_DOCS = (
        "the pizza pizza beer copyright",
        "the pizza burger beer copyright",
        "the the pizza beer beer copyright",
        "the burger beer beer copyright",
        "the coke burger coke copyright",
        "the coke burger burger",
    )
    word_vect = CountVectorizer(analyzer="word")
    char_vect = CountVectorizer(analyzer="char_wb", ngram_range=(3, 3))
    ft = FeatureStacker([("chars", char_vect), ("words", word_vect)])
    ft.fit(JUNK_FOOD_DOCS)
    feature_names = ft.get_feature_names()
    for feat in feature_names:
        assert_true("chars__" in feat or "words__" in feat)
    assert_equal(len(feature_names), 35)
<commit_msg>TST add test for feature weights in feature stacker<commit_after>"""
Test the pipeline module.
"""
import numpy as np
from scipy import sparse

from nose.tools import assert_raises, assert_equal, assert_false, assert_true
from numpy.testing import assert_array_equal, \
        assert_array_almost_equal

from sklearn.base import BaseEstimator, clone
from sklearn.pipeline import Pipeline, FeatureStacker
from sklearn.svm import SVC
from sklearn.linear_model import LogisticRegression
from sklearn.feature_selection import SelectKBest, f_classif
from sklearn.decomposition.pca import PCA, RandomizedPCA
from sklearn.datasets import load_iris
from sklearn.preprocessing import StandardScaler
from sklearn.feature_extraction.text import CountVectorizer

class IncorrectT(BaseEstimator):
    """Small class to test parameter dispatching.
    """

    def __init__(self, a=None, b=None):
        self.a = a
        self.b = b

class T(IncorrectT):

    def fit(self, X, y):
        return self

class TransfT(T):

    def transform(self, X, y=None):
        return X

class FitParamT(BaseEstimator):
    """Mock classifier
    """

    def __init__(self):
        self.successful = False
        pass

    def fit(self, X, y, should_succeed=False):
        self.successful = should_succeed

    def predict(self, X):
        return self.successful

def test_pipeline_init():
    """ Test the various init parameters of the pipeline.
    """
    assert_raises(TypeError, Pipeline)
    # Check that we can't instantiate pipelines with objects without fit
    # method
    pipe = assert_raises(TypeError, Pipeline,
                        [('svc', IncorrectT)])
    # Smoke test with only an estimator
    clf = T()
    pipe = Pipeline([('svc', clf)])
    assert_equal(pipe.get_params(deep=True),
                 dict(svc__a=None, svc__b=None, svc=clf))

    # Check that params are set
    pipe.set_params(svc__a=0.1)
    assert_equal(clf.a, 0.1)
    # Smoke test the repr:
    repr(pipe)

    # Test with two objects
    clf = SVC()
    filter1 = SelectKBest(f_classif)
    pipe = Pipeline([('anova', filter1), ('svc', clf)])

    # Check that params are set
    pipe.set_params(svc__C=0.1)
    assert_equal(clf.C, 0.1)
    # Smoke test the repr:
    repr(pipe)

    # Check that params are not set when naming them wrong
    assert_raises(ValueError, pipe.set_params, anova__C=0.1)

    # Test clone
    pipe2 = clone(pipe)
    assert_false(pipe.named_steps['svc'] is pipe2.named_steps['svc'])

    # Check that appart from estimators, the parameters are the same
    params = pipe.get_params()
    params2 = pipe2.get_params()
    # Remove estimators that where copied
    params.pop('svc')
    params.pop('anova')
    params2.pop('svc')
    params2.pop('anova')
    assert_equal(params, params2)

def test_pipeline_methods_anova():
    """ Test the various methods of the pipeline (anova).
    """
    iris = load_iris()
    X = iris.data
    y = iris.target
    # Test with Anova + LogisticRegression
    clf = LogisticRegression()
    filter1 = SelectKBest(f_classif, k=2)
    pipe = Pipeline([('anova', filter1), ('logistic', clf)])
    pipe.fit(X, y)
    pipe.predict(X)
    pipe.predict_proba(X)
    pipe.predict_log_proba(X)
    pipe.score(X, y)

def test_pipeline_fit_params():
    """Test that the pipeline can take fit parameters
    """
    pipe = Pipeline([('transf', TransfT()), ('clf', FitParamT())])
    pipe.fit(X=None, y=None, clf__should_succeed=True)
    # classifier should return True
    assert_true(pipe.predict(None))
    # and transformer params should not be changed
    assert_true(pipe.named_steps['transf'].a is None)
    assert_true(pipe.named_steps['transf'].b is None)

def test_pipeline_methods_pca_svm():
    """Test the various methods of the pipeline (pca + svm)."""
    iris = load_iris()
    X = iris.data
    y = iris.target
    # Test with PCA + SVC
    clf = SVC(probability=True)
    pca = PCA(n_components='mle', whiten=True)
    pipe = Pipeline([('pca', pca), ('svc', clf)])
    pipe.fit(X, y)
    pipe.predict(X)
    pipe.predict_proba(X)
    pipe.predict_log_proba(X)
    pipe.score(X, y)

def test_pipeline_methods_preprocessing_svm():
    """Test the various methods of the pipeline (preprocessing + svm)."""
    iris = load_iris()
    X = iris.data
    y = iris.target
    n_samples = X.shape[0]
    n_classes = len(np.unique(y))
    scaler = StandardScaler()
    pca = RandomizedPCA(n_components=2, whiten=True)
    clf = SVC(probability=True)

    for preprocessing in [scaler, pca]:
        pipe = Pipeline([('scaler', scaler), ('svc', clf)])
        pipe.fit(X, y)

        # check shapes of various prediction functions
        predict = pipe.predict(X)
        assert_equal(predict.shape, (n_samples,))

        proba = pipe.predict_proba(X)
        assert_equal(proba.shape, (n_samples, n_classes))

        log_proba = pipe.predict_log_proba(X)
        assert_equal(log_proba.shape, (n_samples, n_classes))

        decision_function = pipe.decision_function(X)
        assert_equal(decision_function.shape, (n_samples, n_classes))

        pipe.score(X, y)

def test_feature_stacker():
    # basic sanity check for feature stacker
    iris = load_iris()
    X = iris.data
    X -= X.mean(axis=0)
    y = iris.target
    pca = RandomizedPCA(n_components=2)
    select = SelectKBest(k=1)
    fs = FeatureStacker([("pca", pca), ("select", select)])
    fs.fit(X, y)
    X_transformed = fs.transform(X)
    assert_equal(X_transformed.shape, (X.shape[0], 3))

    # check if it does the expected thing
    assert_array_almost_equal(X_transformed[:, :-1], pca.fit_transform(X))
    assert_array_equal(X_transformed[:, -1],
            select.fit_transform(X, y).ravel())

    # test if it also works for sparse input
    X_sp = sparse.csr_matrix(X)
    X_sp_transformed = fs.fit_transform(X_sp, y)
    assert_array_almost_equal(X_transformed, X_sp_transformed.toarray())

    # test setting parameters
    fs.set_params(select__k=2)
    assert_equal(fs.fit_transform(X, y).shape, (X.shape[0], 4))

def test_feature_stacker_weights():
    # test feature stacker with transformer weights
    iris = load_iris()
    X = iris.data
    y = iris.target
    pca = RandomizedPCA(n_components=2)
    select = SelectKBest(k=1)
    fs = FeatureStacker([("pca", pca), ("select", select)],
            transformer_weights={"pca": 10})
    fs.fit(X, y)
    X_transformed = fs.transform(X)
    # check against expected result
    assert_array_almost_equal(X_transformed[:, :-1], 10 * pca.fit_transform(X))
    assert_array_equal(X_transformed[:, -1],
            select.fit_transform(X, y).ravel())

def test_feature_stacker_feature_names():
    JUNK_FOOD_DOCS = (
        "the pizza pizza beer copyright",
        "the pizza burger beer copyright",
        "the the pizza beer beer copyright",
        "the burger beer beer copyright",
        "the coke burger coke copyright",
        "the coke burger burger",
    )
    word_vect = CountVectorizer(analyzer="word")
    char_vect = CountVectorizer(analyzer="char_wb", ngram_range=(3, 3))
    ft = FeatureStacker([("chars", char_vect), ("words", word_vect)])
    ft.fit(JUNK_FOOD_DOCS)
    feature_names = ft.get_feature_names()
    for feat in feature_names:
        assert_true("chars__" in feat or "words__" in feat)
    assert_equal(len(feature_names), 35)
<|endoftext|>


<commit_before>from functools import partial

import numpy as np

from sklearn.tree import DecisionTreeClassifier
from sklearn.utils.testing import assert_less
from sklearn.utils.testing import assert_raise_message

from skopt import gbrt_minimize
from skopt import forest_minimize
from skopt.benchmarks import bench1
from skopt.benchmarks import bench2
from skopt.benchmarks import bench3
from skopt.benchmarks import bench4
from skopt.benchmarks import branin
from skopt.benchmarks import hart6

MINIMIZERS = [("et", partial(forest_minimize, base_estimator='et')),
              ("rf", partial(forest_minimize, base_estimator='rf')),
              ("gbrt", gbrt_minimize)]

def test_forest_minimize_api():
    # invalid string value
    assert_raise_message(ValueError,
                         "Valid values for the base_estimator parameter",
                         forest_minimize, lambda x: 0., [],
                         base_estimator='abc')

    # not a string nor a regressor
    assert_raise_message(ValueError,
                         "The base_estimator parameter has to either",
                         forest_minimize, lambda x: 0., [],
                         base_estimator=42)

    assert_raise_message(ValueError,
                         "The base_estimator parameter has to either",
                         forest_minimize, lambda x: 0., [],
                         base_estimator=DecisionTreeClassifier())

def check_minimize(minimizer, func, y_opt, dimensions, margin,
                        n_calls, n_random_starts=10, x0=None):
    for n in range(3):
        r = minimizer(
            func, dimensions, n_calls=n_calls, random_state=n,
            n_random_starts=n_random_starts, x0=x0)
        assert_less(r.fun, y_opt + margin)

def test_tree_based_minimize():
    rng = np.random.RandomState(0)
    for name, minimizer in MINIMIZERS:
        yield (check_minimize, minimizer, bench1, 0.,
               [(-2.0, 2.0)], 0.05, 25, 5)

        # This benchmark contains 2 different functions
        # when x < 0 and otherwise.
        # Hence sample uniformly from [-6, 0] and
        # [0, 6] and provide a warm-start.
        X1 = (-6 + 6 * rng.rand(5, 1)).tolist()
        X2 = (6 * rng.rand(5, 1)).tolist()
        X0 = X1 + X2
        yield (check_minimize, minimizer, bench2, -5,
               [(-6.0, 6.0)], 0.05, 100, 0, X0)

        yield (check_minimize, minimizer, bench3, -0.9,
               [(-2.0, 2.0)], 0.05, 25)
        yield (check_minimize, minimizer, bench4, 0.0,
               [("-2", "-1", "0", "1", "2")], 0.05, 10)
        yield (check_minimize, minimizer, hart6, -3.32,
               np.tile((0.0, 1.0), (6, 1)), 1.0, 50)

        if name == "et":
            yield (check_minimize, minimizer, branin, 0.39,
                   [(-5.0, 10.0), (0.0, 15.0)], 0.15, 125)
        else:
            yield (check_minimize, minimizer, branin, 0.39,
                   [(-5.0, 10.0), (0.0, 15.0)], 0.15, 200)
<commit_msg>Seed with points closer to the edge of the space<commit_after>from functools import partial

import numpy as np

from sklearn.tree import DecisionTreeClassifier
from sklearn.utils.testing import assert_less
from sklearn.utils.testing import assert_raise_message

from skopt import gbrt_minimize
from skopt import forest_minimize
from skopt.benchmarks import bench1
from skopt.benchmarks import bench2
from skopt.benchmarks import bench3
from skopt.benchmarks import bench4
from skopt.benchmarks import branin
from skopt.benchmarks import hart6

MINIMIZERS = [("et", partial(forest_minimize, base_estimator='et')),
              ("rf", partial(forest_minimize, base_estimator='rf')),
              ("gbrt", gbrt_minimize)]

def test_forest_minimize_api():
    # invalid string value
    assert_raise_message(ValueError,
                         "Valid values for the base_estimator parameter",
                         forest_minimize, lambda x: 0., [],
                         base_estimator='abc')

    # not a string nor a regressor
    assert_raise_message(ValueError,
                         "The base_estimator parameter has to either",
                         forest_minimize, lambda x: 0., [],
                         base_estimator=42)

    assert_raise_message(ValueError,
                         "The base_estimator parameter has to either",
                         forest_minimize, lambda x: 0., [],
                         base_estimator=DecisionTreeClassifier())

def check_minimize(minimizer, func, y_opt, dimensions, margin,
                        n_calls, n_random_starts=10, x0=None):
    for n in range(3):
        r = minimizer(
            func, dimensions, n_calls=n_calls, random_state=n,
            n_random_starts=n_random_starts, x0=x0)
        assert_less(r.fun, y_opt + margin)

def test_tree_based_minimize():
    for name, minimizer in MINIMIZERS:
        yield (check_minimize, minimizer, bench1, 0.,
               [(-2.0, 2.0)], 0.05, 25, 5)

        # XXX: We supply points at the edge of the search
        # space as an initial point to the minimizer.
        # This makes sure that the RF model can find the minimum even
        # if all the randomly sampled points are one side of the
        # the minimum, since for a decision tree any point greater than
        # max(sampled_points) would give a constant value.
        X0 = [[-5.6], [-5.8], [5.8], [5.6]]
        yield (check_minimize, minimizer, bench2, -5,
               [(-6.0, 6.0)], 0.1, 100, 10, X0)

        yield (check_minimize, minimizer, bench3, -0.9,
               [(-2.0, 2.0)], 0.05, 25)
        yield (check_minimize, minimizer, bench4, 0.0,
               [("-2", "-1", "0", "1", "2")], 0.05, 10)
        yield (check_minimize, minimizer, hart6, -3.32,
               np.tile((0.0, 1.0), (6, 1)), 1.0, 50)

        if name == "et":
            yield (check_minimize, minimizer, branin, 0.39,
                   [(-5.0, 10.0), (0.0, 15.0)], 0.15, 125)
        else:
            yield (check_minimize, minimizer, branin, 0.39,
                   [(-5.0, 10.0), (0.0, 15.0)], 0.15, 200)
<|endoftext|>


<commit_before>from datetime import datetime
from collections import OrderedDict

from sqlalchemy import ForeignKey, Column
from sqlalchemy.orm import relationship, contains_eager
from sqlalchemy.types import Integer, DateTime

from .base import DeclarativeBase
from .session import DBSession
from .auth import User
from .flight import Flight
from .flight_comment import FlightComment
from .follower import Follower

class Event(DeclarativeBase):
    __tablename__ = 'events'

    id = Column(Integer, autoincrement=True, primary_key=True)

    # Notification type

    type = Column(Integer, nullable=False)

    class Type:
        FLIGHT_COMMENT = 1
        FLIGHT = 2
        FOLLOWER = 3

    # Event time

    time = Column(DateTime, nullable=False, default=datetime.utcnow)

    # The user that caused the event

    actor_id = Column(
        Integer, ForeignKey('tg_user.id', ondelete='CASCADE'), nullable=False)
    actor = relationship('User', innerjoin=True)

    # A flight if this event is about a flight

    flight_id = Column(
        Integer, ForeignKey('flights.id', ondelete='CASCADE'))
    flight = relationship('Flight')

    # A flight comment if this event is about a flight comment

    flight_comment_id = Column(
        Integer, ForeignKey('flight_comments.id', ondelete='CASCADE'))
    flight_comment = relationship('FlightComment')

    ##############################

    def __repr__(self):
        return '<Event: id={} type={}>' \
            .format(self.id, self.type).encode('utf-8')

class Notification(DeclarativeBase):
    __tablename__ = 'notifications'

    id = Column(Integer, autoincrement=True, primary_key=True)

    # The event of this notification

    event_id = Column(
        Integer, ForeignKey('events.id', ondelete='CASCADE'), nullable=False)
    event = relationship('Event', innerjoin=True)

    # The recipient of this notification

    recipient_id = Column(
        Integer, ForeignKey('tg_user.id', ondelete='CASCADE'), nullable=False)
    recipient = relationship('User', innerjoin=True)

    # The time that this notification was read by the recipient

    time_read = Column(DateTime)

    ##############################

    def __repr__(self):
        return '<Notification: id={}>' \
            .format(self.id).encode('utf-8')

    ##############################

    @classmethod
    def query_unread(cls, user):
        return cls.query(recipient=user, time_read=None)

    @classmethod
    def count_unread(cls, user):
        return cls.query_unread(user).count()

    ##############################

    def mark_read(self):
        self.time_read = datetime.utcnow()

    @classmethod
    def mark_all_read(cls, user, filter_func=None):
        query = cls.query(recipient=user) \
            .outerjoin(Event) \
            .filter(Event.id == Notification.event_id)

        if filter_func is not None:
            query = filter_func(query)

        query.update(dict(time_read=datetime.utcnow()))

def create_flight_comment_notifications(comment):
    '''
    Create notifications for the owner and pilots of the flight
    '''

    # Create the event
    event = Event(type=Event.Type.FLIGHT_COMMENT,
                  actor=comment.user,
                  flight=comment.flight,
                  flight_comment=comment)

    DBSession.add(event)

    # Create list of potential recipients (using Set to avoid duplicates)
    recipients = {comment.flight.igc_file.owner,
                  comment.flight.pilot,
                  comment.flight.co_pilot}

    # Create notifications for the recipients in the set
    for recipient in recipients:
        # There is no need to notify the user that caused the notification
        if recipient is None or recipient == comment.user:
            continue

        item = Notification(event=event, recipient=recipient)
        DBSession.add(item)

def create_flight_notifications(flight):
    '''
    Create notifications for the followers of the owner and pilots of the flight
    '''

    # Create the event
    event = Event(type=Event.Type.FLIGHT,
                  actor_id=flight.igc_file.owner.id,
                  flight=flight)

    DBSession.add(event)

    # Create list of flight-related users
    senders = {flight.pilot_id, flight.co_pilot_id, flight.igc_file.owner.id}
    senders.discard(None)

    # Request followers/recipients of the flight-related users from the DB
    followers = DBSession.query(Follower.source_id.label('id')) \
                         .filter(Follower.destination_id.in_(senders))

    # Determine the recipients
    recipients = set([follower.id for follower in followers])

    # Don't send notifications to the senders if they follow each other
    recipients.difference_update(senders)

    # Create notifications for the recipients
    for recipient in recipients:
        item = Notification(event=event, recipient_id=recipient)
        DBSession.add(item)

def create_follower_notification(followed, follower):
    '''
    Create notification for the followed pilot about his new follower
    '''

    # Create the event
    event = Event(type=Event.Type.FOLLOWER, actor=follower)
    DBSession.add(event)

    # Create the notification
    item = Notification(event=event, recipient=followed)
    DBSession.add(item)
<commit_msg>model/notification: Simplified create_flight_comment_notifications()<commit_after>from datetime import datetime
from collections import OrderedDict

from sqlalchemy import ForeignKey, Column
from sqlalchemy.orm import relationship, contains_eager
from sqlalchemy.types import Integer, DateTime

from .base import DeclarativeBase
from .session import DBSession
from .auth import User
from .flight import Flight
from .flight_comment import FlightComment
from .follower import Follower

class Event(DeclarativeBase):
    __tablename__ = 'events'

    id = Column(Integer, autoincrement=True, primary_key=True)

    # Notification type

    type = Column(Integer, nullable=False)

    class Type:
        FLIGHT_COMMENT = 1
        FLIGHT = 2
        FOLLOWER = 3

    # Event time

    time = Column(DateTime, nullable=False, default=datetime.utcnow)

    # The user that caused the event

    actor_id = Column(
        Integer, ForeignKey('tg_user.id', ondelete='CASCADE'), nullable=False)
    actor = relationship('User', innerjoin=True)

    # A flight if this event is about a flight

    flight_id = Column(
        Integer, ForeignKey('flights.id', ondelete='CASCADE'))
    flight = relationship('Flight')

    # A flight comment if this event is about a flight comment

    flight_comment_id = Column(
        Integer, ForeignKey('flight_comments.id', ondelete='CASCADE'))
    flight_comment = relationship('FlightComment')

    ##############################

    def __repr__(self):
        return '<Event: id={} type={}>' \
            .format(self.id, self.type).encode('utf-8')

class Notification(DeclarativeBase):
    __tablename__ = 'notifications'

    id = Column(Integer, autoincrement=True, primary_key=True)

    # The event of this notification

    event_id = Column(
        Integer, ForeignKey('events.id', ondelete='CASCADE'), nullable=False)
    event = relationship('Event', innerjoin=True)

    # The recipient of this notification

    recipient_id = Column(
        Integer, ForeignKey('tg_user.id', ondelete='CASCADE'), nullable=False)
    recipient = relationship('User', innerjoin=True)

    # The time that this notification was read by the recipient

    time_read = Column(DateTime)

    ##############################

    def __repr__(self):
        return '<Notification: id={}>' \
            .format(self.id).encode('utf-8')

    ##############################

    @classmethod
    def query_unread(cls, user):
        return cls.query(recipient=user, time_read=None)

    @classmethod
    def count_unread(cls, user):
        return cls.query_unread(user).count()

    ##############################

    def mark_read(self):
        self.time_read = datetime.utcnow()

    @classmethod
    def mark_all_read(cls, user, filter_func=None):
        query = cls.query(recipient=user) \
            .outerjoin(Event) \
            .filter(Event.id == Notification.event_id)

        if filter_func is not None:
            query = filter_func(query)

        query.update(dict(time_read=datetime.utcnow()))

def create_flight_comment_notifications(comment):
    '''
    Create notifications for the owner and pilots of the flight
    '''

    flight = comment.flight
    user = comment.user

    # Create the event
    event = Event(type=Event.Type.FLIGHT_COMMENT,
                  actor=user, flight=flight, flight_comment=comment)

    DBSession.add(event)

    # Create set of potential recipients
    recipients = {flight.igc_file.owner, flight.pilot, flight.co_pilot}
    recipients.discard(None)
    recipients.discard(user)

    # Create notifications for the recipients in the set
    for recipient in recipients:
        item = Notification(event=event, recipient=recipient)
        DBSession.add(item)

def create_flight_notifications(flight):
    '''
    Create notifications for the followers of the owner and pilots of the flight
    '''

    # Create the event
    event = Event(type=Event.Type.FLIGHT,
                  actor_id=flight.igc_file.owner.id,
                  flight=flight)

    DBSession.add(event)

    # Create list of flight-related users
    senders = {flight.pilot_id, flight.co_pilot_id, flight.igc_file.owner.id}
    senders.discard(None)

    # Request followers/recipients of the flight-related users from the DB
    followers = DBSession.query(Follower.source_id.label('id')) \
                         .filter(Follower.destination_id.in_(senders))

    # Determine the recipients
    recipients = set([follower.id for follower in followers])

    # Don't send notifications to the senders if they follow each other
    recipients.difference_update(senders)

    # Create notifications for the recipients
    for recipient in recipients:
        item = Notification(event=event, recipient_id=recipient)
        DBSession.add(item)

def create_follower_notification(followed, follower):
    '''
    Create notification for the followed pilot about his new follower
    '''

    # Create the event
    event = Event(type=Event.Type.FOLLOWER, actor=follower)
    DBSession.add(event)

    # Create the notification
    item = Notification(event=event, recipient=followed)
    DBSession.add(item)
<|endoftext|>


<commit_before>from slm_lab.agent import net
from slm_lab.agent.algorithm import policy_util
from slm_lab.agent.algorithm.sarsa import SARSA
from slm_lab.agent.net import net_util
from slm_lab.lib import logger, util
from slm_lab.lib.decorator import lab_api
import numpy as np
import pydash as ps
import torch

logger = logger.get_logger(__name__)

class VanillaDQN(SARSA):
    '''
    Implementation of a simple DQN algorithm.
    Algorithm:
        1. Collect some examples by acting in the environment and store them in a replay memory
        2. Every K steps sample N examples from replay memory
        3. For each example calculate the target (bootstrapped estimate of the discounted value of the state and action taken), y, using a neural network to approximate the Q function. s' is the next state following the action actually taken.
                y_t = r_t + gamma * argmax_a Q(s_t', a)
        4. For each example calculate the current estimate of the discounted value of the state and action taken
                x_t = Q(s_t, a_t)
        5. Calculate L(x, y) where L is a regression loss (eg. mse)
        6. Calculate the gradient of L with respect to all the parameters in the network and update the network parameters using the gradient
        7. Repeat steps 3 - 6 M times
        8. Repeat steps 2 - 7 Z times
        9. Repeat steps 1 - 8

    For more information on Q-Learning see Sergey Levine's lectures 6 and 7 from CS294-112 Fall 2017
    https://www.youtube.com/playlist?list=PLkFD6_40KJIznC9CDbVTjAF2oyt8_VAe3

    e.g. algorithm_spec
    "algorithm": {
        "name": "VanillaDQN",
        "action_pdtype": "Argmax",
        "action_policy": "epsilon_greedy",
        "explore_var_spec": {
            "name": "linear_decay",
            "start_val": 1.0,
            "end_val": 0.1,
            "start_step": 10,
            "end_step": 1000,
        },
        "gamma": 0.99,
        "training_batch_epoch": 8,
        "training_epoch": 4,
        "training_frequency": 10,
        "training_start_step": 10,
        "normalize_state": true
    }
    '''

    @lab_api
    def init_algorithm_params(self):
        # set default
        util.set_attr(self, dict(
            action_pdtype='Argmax',
            action_policy='epsilon_greedy',
            explore_var_spec=None,
        ))
        util.set_attr(self, self.algorithm_spec, [
            'action_pdtype',
            'action_policy',
            # explore_var is epsilon, tau or etc. depending on the action policy
            # these control the trade off between exploration and exploitaton
            'explore_var_spec',
            'gamma',  # the discount factor
            'training_batch_epoch',  # how many gradient updates per batch
            'training_epoch',  # how many batches to train each time
            'training_frequency',  # how often to train (once a few timesteps)
            'training_start_step',  # how long before starting training
            'normalize_state',
        ])
        super(VanillaDQN, self).init_algorithm_params()

    @lab_api
    def init_nets(self, global_nets=None):
        '''Initialize the neural network used to learn the Q function from the spec'''
        if self.algorithm_spec['name'] == 'VanillaDQN':
            assert all(k not in self.net_spec for k in ['update_type', 'update_frequency', 'polyak_coef']), 'Network update not available for VanillaDQN; use DQN.'
        if global_nets is None:
            in_dim = self.body.state_dim
            out_dim = net_util.get_out_dim(self.body)
            NetClass = getattr(net, self.net_spec['type'])
            self.net = NetClass(self.net_spec, in_dim, out_dim)
            self.net_names = ['net']
        else:
            util.set_attr(self, global_nets)
            self.net_names = list(global_nets.keys())
        self.post_init_nets()

    def calc_q_loss(self, batch):
        '''Compute the Q value loss using predicted and target Q values from the appropriate networks'''
        q_preds = self.net.wrap_eval(batch['states'])
        act_q_preds = q_preds.gather(-1, batch['actions'].long().unsqueeze(-1)).squeeze(-1)
        next_q_preds = self.net.wrap_eval(batch['next_states'])
        # Bellman equation: compute max_q_targets using reward and max estimated Q values (0 if no next_state)
        max_next_q_preds, _ = next_q_preds.max(dim=-1, keepdim=True)
        max_q_targets = batch['rewards'] + self.gamma * (1 - batch['dones']) * max_next_q_preds
        max_q_targets = max_q_targets.detach()
        q_loss = self.net.loss_fn(act_q_preds, max_q_targets)

        # TODO use the same loss_fn but do not reduce yet
        if 'Prioritized' in util.get_class_name(self.body.memory):  # PER
            errors = torch.abs(max_q_targets - act_q_preds.detach())
            self.body.memory.update_priorities(errors)
        return q_loss

    @lab_api
    def act(self, state):
        '''Selects and returns a discrete action for body using the action policy'''
        return super(VanillaDQN, self).act(state)

    @lab_api
    def sample(self):
        '''Samples a batch from memory of size self.memory_spec['batch_size']'''
        batch = self.body.memory.sample()
        if self.normalize_state:
            batch = policy_util.normalize_states_and_next_states(self.body, batch)
        batch = util.to_torch_batch(batch, self.net.device, self.body.memory.is_episodic)
        return batch

    @lab_api
    def train(self):
        '''
        Completes one training step for the agent if it is time to train.
        i.e. the environment timestep is greater than the minimum training timestep and a multiple of the training_frequency.
        Each training step consists of sampling n batches from the agent's memory.
        For each of the batches, the target Q values (q_targets) are computed and a single training step is taken k times
        Otherwise this function does nothing.
        '''
        if util.in_eval_lab_modes():
            self.body.flush()
            return np.nan
        clock = self.body.env.clock
        tick = clock.get(clock.max_tick_unit)
        self.to_train = (tick > self.training_start_step and tick % self.training_frequency == 0)
        if self.to_train == 1:
            total_loss = torch.tensor(0.0, device=self.net.device)
            for _ in range(self.training_epoch):
                batch = self.sample()
                for _ in range(self.training_batch_epoch):
                    loss = self.calc_q_loss(batch)
                    self.net.training_step(loss=loss, lr_clock=clock)
                    total_loss += loss
            loss = total_loss / (self.training_epoch * self.training_batch_epoch)
            # reset
            self.to_train = 0
            self.body.flush()
            logger.debug(f'Trained {self.name} at epi: {clock.epi}, total_t: {clock.total_t}, t: {clock.t}, total_reward so far: {self.body.memory.total_reward}, loss: {loss:g}')
            return loss.item()
        else:
            return np.nan

    @lab_api
    def update(self):
        '''Update the agent after training'''
        return super(VanillaDQN, self).update()

class DQNBase(VanillaDQN):
    '''
    Implementation of the base DQN algorithm.
    The algorithm follows the same general approach as VanillaDQN but is more general since it allows
    for two different networks (through self.net and self.target_net).

    self.net is used to act, and is the network trained.
    self.target_net is used to estimate the maximum value of the Q-function in the next state when calculating the target (see VanillaDQN comments).
    self.target_net is updated periodically to either match self.net (self.net.update_type = "replace") or to be a weighted average of self.net and the previous self.target_net (self.net.update_type = "polyak")
    If desired, self.target_net can be updated slowly, and this can help to stabilize learning.

    It also allows for different nets to be used to select the action in the next state and to evaluate the value of that action through self.online_net and self.eval_net. This can help reduce the tendency of DQN's to overestimate the value of the Q-function. Following this approach leads to the DoubleDQN algorithm.

    Setting all nets to self.net reduces to the VanillaDQN case.
    '''

    @lab_api
    def init_nets(self, global_nets=None):
        '''Initialize networks'''
        if self.algorithm_spec['name'] == 'DQNBase':
            assert all(k not in self.net_spec for k in ['update_type', 'update_frequency', 'polyak_coef']), 'Network update not available for DQNBase; use DQN.'
        if global_nets is None:
            in_dim = self.body.state_dim
            out_dim = net_util.get_out_dim(self.body)
            NetClass = getattr(net, self.net_spec['type'])
            self.net = NetClass(self.net_spec, in_dim, out_dim)
            self.target_net = NetClass(self.net_spec, in_dim, out_dim)
            self.net_names = ['net', 'target_net']
        else:
            util.set_attr(self, global_nets)
            self.net_names = list(global_nets.keys())
        self.post_init_nets()
        self.online_net = self.target_net
        self.eval_net = self.target_net

    def calc_q_loss(self, batch):
        '''Compute the Q value loss using predicted and target Q values from the appropriate networks'''
        q_preds = self.net.wrap_eval(batch['states'])
        act_q_preds = q_preds.gather(-1, batch['actions'].long().unsqueeze(-1)).squeeze(-1)
        # Use online_net to select actions in next state
        online_next_q_preds = self.online_net.wrap_eval(batch['next_states'])
        # Use eval_net to calculate next_q_preds for actions chosen by online_net
        next_q_preds = self.eval_net.wrap_eval(batch['next_states'])
        max_next_q_preds = next_q_preds.gather(-1, online_next_q_preds.argmax(dim=-1, keepdim=True)).squeeze(-1)
        max_q_targets = batch['rewards'] + self.gamma * (1 - batch['dones']) * max_next_q_preds
        max_q_targets = max_q_targets.detach()
        q_loss = self.net.loss_fn(act_q_preds, max_q_targets)

        # TODO use the same loss_fn but do not reduce yet
        if 'Prioritized' in util.get_class_name(self.body.memory):  # PER
            errors = torch.abs(max_q_targets - act_q_preds.detach())
            self.body.memory.update_priorities(errors)
        return q_loss

    def update_nets(self):
        total_t = self.body.env.clock.total_t
        if total_t % self.net.update_frequency == 0:
            if self.net.update_type == 'replace':
                logger.debug('Updating target_net by replacing')
                net_util.copy(self.net, self.target_net)
                self.online_net = self.target_net
                self.eval_net = self.target_net
            elif self.net.update_type == 'polyak':
                logger.debug('Updating net by averaging')
                net_util.polyak_update(self.net, self.target_net, self.net.polyak_coef)
                self.online_net = self.target_net
                self.eval_net = self.target_net
            else:
                raise ValueError('Unknown net.update_type. Should be "replace" or "polyak". Exiting.')

    @lab_api
    def update(self):
        '''Updates self.target_net and the explore variables'''
        self.update_nets()
        return super(DQNBase, self).update()

class DQN(DQNBase):
    '''
    DQN class

    e.g. algorithm_spec
    "algorithm": {
        "name": "DQN",
        "action_pdtype": "Argmax",
        "action_policy": "epsilon_greedy",
        "explore_var_spec": {
            "name": "linear_decay",
            "start_val": 1.0,
            "end_val": 0.1,
            "start_step": 10,
            "end_step": 1000,
        },
        "gamma": 0.99,
        "training_batch_epoch": 8,
        "training_epoch": 4,
        "training_frequency": 10,
        "training_start_step": 10
    }
    '''
    @lab_api
    def init_nets(self, global_nets=None):
        super(DQN, self).init_nets(global_nets)

class DoubleDQN(DQN):
    '''
    Double-DQN (DDQN) class

    e.g. algorithm_spec
    "algorithm": {
        "name": "DDQN",
        "action_pdtype": "Argmax",
        "action_policy": "epsilon_greedy",
        "explore_var_spec": {
            "name": "linear_decay",
            "start_val": 1.0,
            "end_val": 0.1,
            "start_step": 10,
            "end_step": 1000,
        },
        "gamma": 0.99,
        "training_batch_epoch": 8,
        "training_epoch": 4,
        "training_frequency": 10,
        "training_start_step": 10
    }
    '''
    @lab_api
    def init_nets(self, global_nets=None):
        super(DoubleDQN, self).init_nets(global_nets)
        self.online_net = self.net
        self.eval_net = self.target_net

    def update_nets(self):
        res = super(DoubleDQN, self).update_nets()
        total_t = self.body.env.clock.total_t
        if self.net.update_type == 'replace':
            if total_t % self.net.update_frequency == 0:
                self.online_net = self.net
                self.eval_net = self.target_net
        elif self.net.update_type == 'polyak':
            self.online_net = self.net
            self.eval_net = self.target_net
<commit_msg>remove redundant net ref setting in DQN and DDQN<commit_after>from slm_lab.agent import net
from slm_lab.agent.algorithm import policy_util
from slm_lab.agent.algorithm.sarsa import SARSA
from slm_lab.agent.net import net_util
from slm_lab.lib import logger, util
from slm_lab.lib.decorator import lab_api
import numpy as np
import pydash as ps
import torch

logger = logger.get_logger(__name__)

class VanillaDQN(SARSA):
    '''
    Implementation of a simple DQN algorithm.
    Algorithm:
        1. Collect some examples by acting in the environment and store them in a replay memory
        2. Every K steps sample N examples from replay memory
        3. For each example calculate the target (bootstrapped estimate of the discounted value of the state and action taken), y, using a neural network to approximate the Q function. s' is the next state following the action actually taken.
                y_t = r_t + gamma * argmax_a Q(s_t', a)
        4. For each example calculate the current estimate of the discounted value of the state and action taken
                x_t = Q(s_t, a_t)
        5. Calculate L(x, y) where L is a regression loss (eg. mse)
        6. Calculate the gradient of L with respect to all the parameters in the network and update the network parameters using the gradient
        7. Repeat steps 3 - 6 M times
        8. Repeat steps 2 - 7 Z times
        9. Repeat steps 1 - 8

    For more information on Q-Learning see Sergey Levine's lectures 6 and 7 from CS294-112 Fall 2017
    https://www.youtube.com/playlist?list=PLkFD6_40KJIznC9CDbVTjAF2oyt8_VAe3

    e.g. algorithm_spec
    "algorithm": {
        "name": "VanillaDQN",
        "action_pdtype": "Argmax",
        "action_policy": "epsilon_greedy",
        "explore_var_spec": {
            "name": "linear_decay",
            "start_val": 1.0,
            "end_val": 0.1,
            "start_step": 10,
            "end_step": 1000,
        },
        "gamma": 0.99,
        "training_batch_epoch": 8,
        "training_epoch": 4,
        "training_frequency": 10,
        "training_start_step": 10,
        "normalize_state": true
    }
    '''

    @lab_api
    def init_algorithm_params(self):
        # set default
        util.set_attr(self, dict(
            action_pdtype='Argmax',
            action_policy='epsilon_greedy',
            explore_var_spec=None,
        ))
        util.set_attr(self, self.algorithm_spec, [
            'action_pdtype',
            'action_policy',
            # explore_var is epsilon, tau or etc. depending on the action policy
            # these control the trade off between exploration and exploitaton
            'explore_var_spec',
            'gamma',  # the discount factor
            'training_batch_epoch',  # how many gradient updates per batch
            'training_epoch',  # how many batches to train each time
            'training_frequency',  # how often to train (once a few timesteps)
            'training_start_step',  # how long before starting training
            'normalize_state',
        ])
        super(VanillaDQN, self).init_algorithm_params()

    @lab_api
    def init_nets(self, global_nets=None):
        '''Initialize the neural network used to learn the Q function from the spec'''
        if self.algorithm_spec['name'] == 'VanillaDQN':
            assert all(k not in self.net_spec for k in ['update_type', 'update_frequency', 'polyak_coef']), 'Network update not available for VanillaDQN; use DQN.'
        if global_nets is None:
            in_dim = self.body.state_dim
            out_dim = net_util.get_out_dim(self.body)
            NetClass = getattr(net, self.net_spec['type'])
            self.net = NetClass(self.net_spec, in_dim, out_dim)
            self.net_names = ['net']
        else:
            util.set_attr(self, global_nets)
            self.net_names = list(global_nets.keys())
        self.post_init_nets()

    def calc_q_loss(self, batch):
        '''Compute the Q value loss using predicted and target Q values from the appropriate networks'''
        q_preds = self.net.wrap_eval(batch['states'])
        act_q_preds = q_preds.gather(-1, batch['actions'].long().unsqueeze(-1)).squeeze(-1)
        next_q_preds = self.net.wrap_eval(batch['next_states'])
        # Bellman equation: compute max_q_targets using reward and max estimated Q values (0 if no next_state)
        max_next_q_preds, _ = next_q_preds.max(dim=-1, keepdim=True)
        max_q_targets = batch['rewards'] + self.gamma * (1 - batch['dones']) * max_next_q_preds
        max_q_targets = max_q_targets.detach()
        q_loss = self.net.loss_fn(act_q_preds, max_q_targets)

        # TODO use the same loss_fn but do not reduce yet
        if 'Prioritized' in util.get_class_name(self.body.memory):  # PER
            errors = torch.abs(max_q_targets - act_q_preds.detach())
            self.body.memory.update_priorities(errors)
        return q_loss

    @lab_api
    def act(self, state):
        '''Selects and returns a discrete action for body using the action policy'''
        return super(VanillaDQN, self).act(state)

    @lab_api
    def sample(self):
        '''Samples a batch from memory of size self.memory_spec['batch_size']'''
        batch = self.body.memory.sample()
        if self.normalize_state:
            batch = policy_util.normalize_states_and_next_states(self.body, batch)
        batch = util.to_torch_batch(batch, self.net.device, self.body.memory.is_episodic)
        return batch

    @lab_api
    def train(self):
        '''
        Completes one training step for the agent if it is time to train.
        i.e. the environment timestep is greater than the minimum training timestep and a multiple of the training_frequency.
        Each training step consists of sampling n batches from the agent's memory.
        For each of the batches, the target Q values (q_targets) are computed and a single training step is taken k times
        Otherwise this function does nothing.
        '''
        if util.in_eval_lab_modes():
            self.body.flush()
            return np.nan
        clock = self.body.env.clock
        tick = clock.get(clock.max_tick_unit)
        self.to_train = (tick > self.training_start_step and tick % self.training_frequency == 0)
        if self.to_train == 1:
            total_loss = torch.tensor(0.0, device=self.net.device)
            for _ in range(self.training_epoch):
                batch = self.sample()
                for _ in range(self.training_batch_epoch):
                    loss = self.calc_q_loss(batch)
                    self.net.training_step(loss=loss, lr_clock=clock)
                    total_loss += loss
            loss = total_loss / (self.training_epoch * self.training_batch_epoch)
            # reset
            self.to_train = 0
            self.body.flush()
            logger.debug(f'Trained {self.name} at epi: {clock.epi}, total_t: {clock.total_t}, t: {clock.t}, total_reward so far: {self.body.memory.total_reward}, loss: {loss:g}')
            return loss.item()
        else:
            return np.nan

    @lab_api
    def update(self):
        '''Update the agent after training'''
        return super(VanillaDQN, self).update()

class DQNBase(VanillaDQN):
    '''
    Implementation of the base DQN algorithm.
    The algorithm follows the same general approach as VanillaDQN but is more general since it allows
    for two different networks (through self.net and self.target_net).

    self.net is used to act, and is the network trained.
    self.target_net is used to estimate the maximum value of the Q-function in the next state when calculating the target (see VanillaDQN comments).
    self.target_net is updated periodically to either match self.net (self.net.update_type = "replace") or to be a weighted average of self.net and the previous self.target_net (self.net.update_type = "polyak")
    If desired, self.target_net can be updated slowly, and this can help to stabilize learning.

    It also allows for different nets to be used to select the action in the next state and to evaluate the value of that action through self.online_net and self.eval_net. This can help reduce the tendency of DQN's to overestimate the value of the Q-function. Following this approach leads to the DoubleDQN algorithm.

    Setting all nets to self.net reduces to the VanillaDQN case.
    '''

    @lab_api
    def init_nets(self, global_nets=None):
        '''Initialize networks'''
        if self.algorithm_spec['name'] == 'DQNBase':
            assert all(k not in self.net_spec for k in ['update_type', 'update_frequency', 'polyak_coef']), 'Network update not available for DQNBase; use DQN.'
        if global_nets is None:
            in_dim = self.body.state_dim
            out_dim = net_util.get_out_dim(self.body)
            NetClass = getattr(net, self.net_spec['type'])
            self.net = NetClass(self.net_spec, in_dim, out_dim)
            self.target_net = NetClass(self.net_spec, in_dim, out_dim)
            self.net_names = ['net', 'target_net']
        else:
            util.set_attr(self, global_nets)
            self.net_names = list(global_nets.keys())
        self.post_init_nets()
        self.online_net = self.target_net
        self.eval_net = self.target_net

    def calc_q_loss(self, batch):
        '''Compute the Q value loss using predicted and target Q values from the appropriate networks'''
        q_preds = self.net.wrap_eval(batch['states'])
        act_q_preds = q_preds.gather(-1, batch['actions'].long().unsqueeze(-1)).squeeze(-1)
        # Use online_net to select actions in next state
        online_next_q_preds = self.online_net.wrap_eval(batch['next_states'])
        # Use eval_net to calculate next_q_preds for actions chosen by online_net
        next_q_preds = self.eval_net.wrap_eval(batch['next_states'])
        max_next_q_preds = next_q_preds.gather(-1, online_next_q_preds.argmax(dim=-1, keepdim=True)).squeeze(-1)
        max_q_targets = batch['rewards'] + self.gamma * (1 - batch['dones']) * max_next_q_preds
        max_q_targets = max_q_targets.detach()
        q_loss = self.net.loss_fn(act_q_preds, max_q_targets)

        # TODO use the same loss_fn but do not reduce yet
        if 'Prioritized' in util.get_class_name(self.body.memory):  # PER
            errors = torch.abs(max_q_targets - act_q_preds.detach())
            self.body.memory.update_priorities(errors)
        return q_loss

    def update_nets(self):
        total_t = self.body.env.clock.total_t
        if total_t % self.net.update_frequency == 0:
            if self.net.update_type == 'replace':
                logger.debug('Updating target_net by replacing')
                net_util.copy(self.net, self.target_net)
            elif self.net.update_type == 'polyak':
                logger.debug('Updating net by averaging')
                net_util.polyak_update(self.net, self.target_net, self.net.polyak_coef)
            else:
                raise ValueError('Unknown net.update_type. Should be "replace" or "polyak". Exiting.')

    @lab_api
    def update(self):
        '''Updates self.target_net and the explore variables'''
        self.update_nets()
        return super(DQNBase, self).update()

class DQN(DQNBase):
    '''
    DQN class

    e.g. algorithm_spec
    "algorithm": {
        "name": "DQN",
        "action_pdtype": "Argmax",
        "action_policy": "epsilon_greedy",
        "explore_var_spec": {
            "name": "linear_decay",
            "start_val": 1.0,
            "end_val": 0.1,
            "start_step": 10,
            "end_step": 1000,
        },
        "gamma": 0.99,
        "training_batch_epoch": 8,
        "training_epoch": 4,
        "training_frequency": 10,
        "training_start_step": 10
    }
    '''
    @lab_api
    def init_nets(self, global_nets=None):
        super(DQN, self).init_nets(global_nets)

class DoubleDQN(DQN):
    '''
    Double-DQN (DDQN) class

    e.g. algorithm_spec
    "algorithm": {
        "name": "DDQN",
        "action_pdtype": "Argmax",
        "action_policy": "epsilon_greedy",
        "explore_var_spec": {
            "name": "linear_decay",
            "start_val": 1.0,
            "end_val": 0.1,
            "start_step": 10,
            "end_step": 1000,
        },
        "gamma": 0.99,
        "training_batch_epoch": 8,
        "training_epoch": 4,
        "training_frequency": 10,
        "training_start_step": 10
    }
    '''
    @lab_api
    def init_nets(self, global_nets=None):
        super(DoubleDQN, self).init_nets(global_nets)
        self.online_net = self.net
        self.eval_net = self.target_net
<|endoftext|>