<commit_before>import mschap
import hashlib
import random

SHSpad1 = \
    "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00" + \
    "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00" + \
    "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00" + \
    "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"

SHSpad2 = \
    "\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2" + \
    "\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2" + \
    "\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2" + \
    "\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2"

Magic1 = \
    "\x54\x68\x69\x73\x20\x69\x73\x20\x74" + \
    "\x68\x65\x20\x4d\x50\x50\x45\x20\x4d" + \
    "\x61\x73\x74\x65\x72\x20\x4b\x65\x79"

Magic2 = \
    "\x4f\x6e\x20\x74\x68\x65\x20\x63\x6c\x69" + \
    "\x65\x6e\x74\x20\x73\x69\x64\x65\x2c\x20" + \
    "\x74\x68\x69\x73\x20\x69\x73\x20\x74\x68" + \
    "\x65\x20\x73\x65\x6e\x64\x20\x6b\x65\x79" + \
    "\x3b\x20\x6f\x6e\x20\x74\x68\x65\x20\x73" + \
    "\x65\x72\x76\x65\x72\x20\x73\x69\x64\x65" + \
    "\x2c\x20\x69\x74\x20\x69\x73\x20\x74\x68" + \
    "\x65\x20\x72\x65\x63\x65\x69\x76\x65\x20" + \
    "\x6b\x65\x79\x2e"

Magic3 = \
    "\x4f\x6e\x20\x74\x68\x65\x20\x63\x6c\x69" + \
    "\x65\x6e\x74\x20\x73\x69\x64\x65\x2c\x20" + \
    "\x74\x68\x69\x73\x20\x69\x73\x20\x74\x68" + \
    "\x65\x20\x72\x65\x63\x65\x69\x76\x65\x20" + \
    "\x6b\x65\x79\x3b\x20\x6f\x6e\x20\x74\x68" + \
    "\x65\x20\x73\x65\x72\x76\x65\x72\x20\x73" + \
    "\x69\x64\x65\x2c\x20\x69\x74\x20\x69\x73" + \
    "\x20\x74\x68\x65\x20\x73\x65\x6e\x64\x20" + \
    "\x6b\x65\x79\x2e"

def mppe_chap2_gen_keys(password, nt_response):
    """
3.3.  Generating 128-bit Session Keys

   When used in conjunction with MS-CHAP-2 authentication, the initial
   MPPE session keys are derived from the peer's Windows NT password.

   The first step is to obfuscate the peer's password using
   NtPasswordHash() function as described in [8].

      NtPasswordHash(Password, PasswordHash)

   The first 16 octets of the result are then hashed again using the MD4
   algorithm.

      PasswordHashHash = md4(PasswordHash)

   The first 16 octets of this second hash are used together with the
   NT-Response field from the MS-CHAP-2 Response packet [8] as the basis
   for the master session key:

      GetMasterKey(PasswordHashHash, NtResponse, MasterKey)

   Once the master key has been generated, it is used to derive two
   128-bit master session keys, one for sending and one for receiving:

GetAsymmetricStartKey(MasterKey, MasterSendKey, 16, TRUE, TRUE)
GetAsymmetricStartKey(MasterKey, MasterReceiveKey, 16, FALSE, TRUE)

   The master session keys are never used to encrypt or decrypt data;
   they are only used in the derivation of transient session keys.  The
   initial transient session keys are obtained by calling the function
   GetNewKeyFromSHA() (described in [3]):

GetNewKeyFromSHA(MasterSendKey, MasterSendKey, 16, SendSessionKey)
GetNewKeyFromSHA(MasterReceiveKey, MasterReceiveKey, 16,
                                                ReceiveSessionKey)

   Finally, the RC4 tables are initialized using the new session keys:

      rc4_key(SendRC4key, 16, SendSessionKey)
      rc4_key(ReceiveRC4key, 16, ReceiveSessionKey)
    """
    password_hash = mschap.nt_password_hash(password)
    password_hash_hash = mschap.hash_nt_password_hash(password_hash)
    master_key = get_master_key(password_hash_hash, nt_response)
    master_send_key = get_asymetric_start_key(master_key, 16, True, True)
    master_recv_key = get_asymetric_start_key(master_key, 16, False, True)
    return master_send_key, master_recv_key

def get_master_key(password_hash_hash, nt_response):
    """
   GetMasterKey(
   IN  16-octet  PasswordHashHash,
   IN  24-octet  NTResponse,
   OUT 16-octet  MasterKey )
   {
      20-octet Digest

      ZeroMemory(Digest, sizeof(Digest));

      /*
       * SHSInit(), SHSUpdate() and SHSFinal()
       * are an implementation of the Secure Hash Standard [7].
       */

      SHSInit(Context);
      SHSUpdate(Context, PasswordHashHash, 16);
      SHSUpdate(Context, NTResponse, 24);
      SHSUpdate(Context, Magic1, 27);
      SHSFinal(Context, Digest);

      MoveMemory(MasterKey, Digest, 16);
   }

    """
    sha_hash = hashlib.sha1()
    sha_hash.update(password_hash_hash)
    sha_hash.update(nt_response)
    sha_hash.update(Magic1)
    return sha_hash.digest()[:16]

def get_asymetric_start_key(master_key, session_key_length, is_send, is_server):
    """

VOID
   GetAsymetricStartKey(
   IN   16-octet      MasterKey,
   OUT  8-to-16 octet SessionKey,
   IN   INTEGER       SessionKeyLength,
   IN   BOOLEAN       IsSend,
   IN   BOOLEAN       IsServer )
   {

      20-octet Digest;

      ZeroMemory(Digest, 20);

      if (IsSend) {
         if (IsServer) {
            s = Magic3
         } else {
            s = Magic2
         }
      } else {
         if (IsServer) {

            s = Magic2
         } else {
            s = Magic3
         }
      }

      /*
       * SHSInit(), SHSUpdate() and SHSFinal()
       * are an implementation of the Secure Hash Standard [7].
       */

      SHSInit(Context);
      SHSUpdate(Context, MasterKey, 16);
      SHSUpdate(Context, SHSpad1, 40);
      SHSUpdate(Context, s, 84);
      SHSUpdate(Context, SHSpad2, 40);
      SHSFinal(Context, Digest);

      MoveMemory(SessionKey, Digest, SessionKeyLength);
   }
    """
    if is_send:
        if is_server:
            s = Magic3
        else:
            s = Magic2
    else:
        if is_server:
            s = Magic2
        else:
            s = Magic3
    sha_hash = hashlib.sha1()
    sha_hash.update(master_key)
    sha_hash.update(SHSpad1)
    sha_hash.update(s)
    sha_hash.update(SHSpad2)
    return sha_hash.digest()[:session_key_length]

def create_plain_text(key):
    key_len = len(key)
    while (len(key) + 1) % 16: key += "\000"
    return chr(key_len) + key

def create_salts():
    send_salt = create_salt()
    recv_salt = create_salt()
    while send_salt == recv_salt: recv_salt = create_salt()
    return (send_salt, recv_salt)

def create_salt():
    return chr(128 + random.randrange(0, 128)) + chr(random.randrange(0, 256))

def gen_radius_encrypt_keys(send_key, recv_key, secret, request_authenticator):
    send_salt, recv_salt = create_salts()
    _send_key = radius_encrypt_keys(
        create_plain_text(send_key),
        secret,
        request_authenticator,
        send_salt
    )
    _recv_key = radius_encrypt_keys(
        create_plain_text(recv_key),
        secret,
        request_authenticator,
        recv_salt
    )

    return _send_key, _recv_key

def radius_encrypt_keys(plain_text, secret, request_authenticator, salt):
    """
  Construct a plaintext version of the String field by concate-
         nating the Key-Length and Key sub-fields.  If necessary, pad
         the resulting string until its length (in octets) is an even
         multiple of 16.  It is recommended that zero octets (0x00) be
         used for padding.  Call this plaintext P.

         Call the shared secret S, the pseudo-random 128-bit Request
         Authenticator (from the corresponding Access-Request packet) R,
         and the contents of the Salt field A.  Break P into 16 octet
         chunks p(1), p(2)...p(i), where i = len(P)/16.  Call the
         ciphertext blocks c(1), c(2)...c(i) and the final ciphertext C.
         Intermediate values b(1), b(2)...c(i) are required.  Encryption
         is performed in the following manner ('+' indicates
         concatenation):

      b(1) = MD5(S + R + A)    c(1) = p(1) xor b(1)   C = c(1)
      b(2) = MD5(S + c(1))     c(2) = p(2) xor b(2)   C = C + c(2)
                  .                      .
                  .                      .
                  .                      .
      b(i) = MD5(S + c(i-1))   c(i) = p(i) xor b(i)   C = C + c(i)

      The   resulting   encrypted   String   field    will    contain
      c(1)+c(2)+...+c(i).
    """
    i = len(plain_text) / 16
    b = hashlib.new("md5", secret + request_authenticator + salt).digest()
    c = xor(plain_text[:16], b)
    result = c
    for x in range(1, i):
        b = hashlib.new("md5", secret + c).digest()
        c = xor(plain_text[x * 16:(x + 1) * 16], b)
        result += c
    return result

def xor(str1, str2):
    return ''.join(map(lambda s1, s2: chr(ord(s1) ^ ord(s2)), str1, str2))
<commit_msg>mschap up<commit_after>import mschap
import hashlib
import random

SHSpad1 = \
    "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00" + \
    "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00" + \
    "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00" + \
    "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"

SHSpad2 = \
    "\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2" + \
    "\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2" + \
    "\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2" + \
    "\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2\xf2"

Magic1 = \
    "\x54\x68\x69\x73\x20\x69\x73\x20\x74" + \
    "\x68\x65\x20\x4d\x50\x50\x45\x20\x4d" + \
    "\x61\x73\x74\x65\x72\x20\x4b\x65\x79"

Magic2 = \
    "\x4f\x6e\x20\x74\x68\x65\x20\x63\x6c\x69" + \
    "\x65\x6e\x74\x20\x73\x69\x64\x65\x2c\x20" + \
    "\x74\x68\x69\x73\x20\x69\x73\x20\x74\x68" + \
    "\x65\x20\x73\x65\x6e\x64\x20\x6b\x65\x79" + \
    "\x3b\x20\x6f\x6e\x20\x74\x68\x65\x20\x73" + \
    "\x65\x72\x76\x65\x72\x20\x73\x69\x64\x65" + \
    "\x2c\x20\x69\x74\x20\x69\x73\x20\x74\x68" + \
    "\x65\x20\x72\x65\x63\x65\x69\x76\x65\x20" + \
    "\x6b\x65\x79\x2e"

Magic3 = \
    "\x4f\x6e\x20\x74\x68\x65\x20\x63\x6c\x69" + \
    "\x65\x6e\x74\x20\x73\x69\x64\x65\x2c\x20" + \
    "\x74\x68\x69\x73\x20\x69\x73\x20\x74\x68" + \
    "\x65\x20\x72\x65\x63\x65\x69\x76\x65\x20" + \
    "\x6b\x65\x79\x3b\x20\x6f\x6e\x20\x74\x68" + \
    "\x65\x20\x73\x65\x72\x76\x65\x72\x20\x73" + \
    "\x69\x64\x65\x2c\x20\x69\x74\x20\x69\x73" + \
    "\x20\x74\x68\x65\x20\x73\x65\x6e\x64\x20" + \
    "\x6b\x65\x79\x2e"

def mppe_chap2_gen_keys(password, nt_response):
    """
3.3.  Generating 128-bit Session Keys

   When used in conjunction with MS-CHAP-2 authentication, the initial
   MPPE session keys are derived from the peer's Windows NT password.

   The first step is to obfuscate the peer's password using
   NtPasswordHash() function as described in [8].

      NtPasswordHash(Password, PasswordHash)

   The first 16 octets of the result are then hashed again using the MD4
   algorithm.

      PasswordHashHash = md4(PasswordHash)

   The first 16 octets of this second hash are used together with the
   NT-Response field from the MS-CHAP-2 Response packet [8] as the basis
   for the master session key:

      GetMasterKey(PasswordHashHash, NtResponse, MasterKey)

   Once the master key has been generated, it is used to derive two
   128-bit master session keys, one for sending and one for receiving:

GetAsymmetricStartKey(MasterKey, MasterSendKey, 16, TRUE, TRUE)
GetAsymmetricStartKey(MasterKey, MasterReceiveKey, 16, FALSE, TRUE)

   The master session keys are never used to encrypt or decrypt data;
   they are only used in the derivation of transient session keys.  The
   initial transient session keys are obtained by calling the function
   GetNewKeyFromSHA() (described in [3]):

GetNewKeyFromSHA(MasterSendKey, MasterSendKey, 16, SendSessionKey)
GetNewKeyFromSHA(MasterReceiveKey, MasterReceiveKey, 16,
                                                ReceiveSessionKey)

   Finally, the RC4 tables are initialized using the new session keys:

      rc4_key(SendRC4key, 16, SendSessionKey)
      rc4_key(ReceiveRC4key, 16, ReceiveSessionKey)
    """
    password_hash = mschap.nt_password_hash(password)
    password_hash_hash = mschap.hash_nt_password_hash(password_hash)
    master_key = get_master_key(password_hash_hash, nt_response)
    master_send_key = get_asymetric_start_key(master_key, 16, True, True)
    master_recv_key = get_asymetric_start_key(master_key, 16, False, True)
    return master_send_key, master_recv_key

def get_master_key(password_hash_hash, nt_response):
    """
   GetMasterKey(
   IN  16-octet  PasswordHashHash,
   IN  24-octet  NTResponse,
   OUT 16-octet  MasterKey )
   {
      20-octet Digest

      ZeroMemory(Digest, sizeof(Digest));

      /*
       * SHSInit(), SHSUpdate() and SHSFinal()
       * are an implementation of the Secure Hash Standard [7].
       */

      SHSInit(Context);
      SHSUpdate(Context, PasswordHashHash, 16);
      SHSUpdate(Context, NTResponse, 24);
      SHSUpdate(Context, Magic1, 27);
      SHSFinal(Context, Digest);

      MoveMemory(MasterKey, Digest, 16);
   }

    """
    sha_hash = hashlib.sha1()
    sha_hash.update(password_hash_hash)
    sha_hash.update(nt_response)
    sha_hash.update(Magic1)
    return sha_hash.digest()[:16]

def get_asymetric_start_key(master_key, session_key_length, is_send, is_server):
    """

VOID
   GetAsymetricStartKey(
   IN   16-octet      MasterKey,
   OUT  8-to-16 octet SessionKey,
   IN   INTEGER       SessionKeyLength,
   IN   BOOLEAN       IsSend,
   IN   BOOLEAN       IsServer )
   {

      20-octet Digest;

      ZeroMemory(Digest, 20);

      if (IsSend) {
         if (IsServer) {
            s = Magic3
         } else {
            s = Magic2
         }
      } else {
         if (IsServer) {

            s = Magic2
         } else {
            s = Magic3
         }
      }

      /*
       * SHSInit(), SHSUpdate() and SHSFinal()
       * are an implementation of the Secure Hash Standard [7].
       */

      SHSInit(Context);
      SHSUpdate(Context, MasterKey, 16);
      SHSUpdate(Context, SHSpad1, 40);
      SHSUpdate(Context, s, 84);
      SHSUpdate(Context, SHSpad2, 40);
      SHSFinal(Context, Digest);

      MoveMemory(SessionKey, Digest, SessionKeyLength);
   }
    """
    if is_send:
        if is_server:
            s = Magic3
        else:
            s = Magic2
    else:
        if is_server:
            s = Magic2
        else:
            s = Magic3
    sha_hash = hashlib.sha1()
    sha_hash.update(master_key)
    sha_hash.update(SHSpad1)
    sha_hash.update(s)
    sha_hash.update(SHSpad2)
    return sha_hash.digest()[:session_key_length]

def create_plain_text(key):
    key_len = len(key)
    while (len(key) + 1) % 16: key += "\000"
    return chr(key_len) + key

def create_salts():
    send_salt = create_salt()
    recv_salt = create_salt()
    while send_salt == recv_salt: recv_salt = create_salt()
    return (send_salt, recv_salt)

def create_salt():
    return chr(128 + random.randrange(0, 128)) + chr(random.randrange(0, 256))

def gen_radius_encrypt_keys(send_key, recv_key, secret, request_authenticator):
    send_salt, recv_salt = create_salts()
    _send_key = send_salt + radius_encrypt_keys(
        create_plain_text(send_key),
        secret,
        request_authenticator,
        send_salt
    )
    _recv_key = recv_salt + radius_encrypt_keys(
        create_plain_text(recv_key),
        secret,
        request_authenticator,
        recv_salt
    )

    return _send_key, _recv_key

def radius_encrypt_keys(plain_text, secret, request_authenticator, salt):
    """
  Construct a plaintext version of the String field by concate-
         nating the Key-Length and Key sub-fields.  If necessary, pad
         the resulting string until its length (in octets) is an even
         multiple of 16.  It is recommended that zero octets (0x00) be
         used for padding.  Call this plaintext P.

         Call the shared secret S, the pseudo-random 128-bit Request
         Authenticator (from the corresponding Access-Request packet) R,
         and the contents of the Salt field A.  Break P into 16 octet
         chunks p(1), p(2)...p(i), where i = len(P)/16.  Call the
         ciphertext blocks c(1), c(2)...c(i) and the final ciphertext C.
         Intermediate values b(1), b(2)...c(i) are required.  Encryption
         is performed in the following manner ('+' indicates
         concatenation):

      b(1) = MD5(S + R + A)    c(1) = p(1) xor b(1)   C = c(1)
      b(2) = MD5(S + c(1))     c(2) = p(2) xor b(2)   C = C + c(2)
                  .                      .
                  .                      .
                  .                      .
      b(i) = MD5(S + c(i-1))   c(i) = p(i) xor b(i)   C = C + c(i)

      The   resulting   encrypted   String   field    will    contain
      c(1)+c(2)+...+c(i).
    """
    i = len(plain_text) / 16
    b = hashlib.new("md5", secret + request_authenticator + salt).digest()
    c = xor(plain_text[:16], b)
    result = c
    for x in range(1, i):
        b = hashlib.new("md5", secret + c).digest()
        c = xor(plain_text[x * 16:(x + 1) * 16], b)
        result += c
    return result

def xor(str1, str2):
    return ''.join(map(lambda s1, s2: chr(ord(s1) ^ ord(s2)), str1, str2))
<|endoftext|>


<commit_before>import json
import os, sys
from shutil import copyfile

from datetime import datetime
import inspect
import tensorflow as tf

if 'SUMO_HOME' in os.environ:
    tools = os.path.join(os.environ['SUMO_HOME'], 'tools')
    sys.path.append(tools)
else:
    sys.exit("please declare environment variable 'SUMO_HOME'")

import gym

from baselines import deepq
import Traci_1_cross_env.Traci_1_cross_env
import Traci_2_cross_env.Traci_2_cross_env
import Traci_3_cross_env.Traci_3_cross_env
from baselines import logger, logger_utils
from BaseTraciEnv import BaseTraciEnv
from pathlib import Path

def callback(lcl, glb):
    # stop training if reward exceeds 199
    is_solved = lcl['t'] > 100 and sum(lcl['episode_rewards'][-101:-1]) / 100 >= 199
    return is_solved

def train_and_log(environment="Traci_3_cross_env-v0",
                  car_chances=1000,
                  reward_function=BaseTraciEnv.reward_average_accumulated_wait_time,
                  lr=1e-3,
                  max_timesteps=int(1e5),
                  buffer_size=10000,
                  exploration_fraction=0.5,
                  explore_final_eps=0.02,
                  train_freq=100,
                  batch_size=32,
                  checkpoint_freq=int(10000),
                  learning_starts=10000,
                  gamma=0.9,
                  target_network_update_freq=500,
                  start_car_probabilities=[0.25, 0.05],  # [0.1,0.1,0.1,0.1,0.1,0.1,0.1], #For traci_3_cross: Bigroad_spawn_prob,Smallroad_spawn_prob
                  end_car_probabilities=[1.0,0.1],
                  num_steps_from_start_car_probs_to_end_car_probs=1e5,
                  prioritized_replay=False,
                  prioritized_replay_alpha=0.6,
                  prioritized_replay_beta0=0.4,
                  prioritized_replay_beta_iters=None,
                  prioritized_replay_eps=1e-6,
                  num_cpu=4,
                  param_noise=False,
                  state_use_queue_length=True,
                  state_use_tl_state=True,
                  state_use_time_since_tl_change=True,
                  state_use_avg_speed=False,
                  hidden_layers=[8, 8, 8],
                  num_actions_pr_trafficlight=3):
    # Print call values
    frame = inspect.currentframe()
    args, _, _, values = inspect.getargvalues(frame)

    call_params_string_array = ['function name "%s"' % inspect.getframeinfo(frame)[2]]
    for i in args:
        call_params_string_array.append("    %s = %s" % (i, values[i]))

    # Setup path of logging, name of environment and save the current arguments (this script)
    log_dir = [os.path.join(str(Path.home()), "Desktop"), environment]
    logger_path = logger_utils.path_with_date(log_dir[0], log_dir[1])

    # Create environment and initialize
    env = gym.make(log_dir[1])
    env.configure_traci(num_car_chances=car_chances,
                        start_car_probabilities=start_car_probabilities,
                        end_car_probabilities=end_car_probabilities,
                        num_steps_from_start_car_probs_to_end_car_probs=num_steps_from_start_car_probs_to_end_car_probs,
                        reward_func=reward_function,
                        state_contain_num_cars_in_queue_history=state_use_queue_length,
                        state_contain_avg_speed_between_detectors_history=state_use_avg_speed,
                        state_contain_time_since_tl_change=state_use_time_since_tl_change,
                        state_contain_tl_state_history=state_use_tl_state,
                        num_actions_pr_trafficlight=num_actions_pr_trafficlight)
    #env.render()

    # Initialize logger
    logger.reset()
    logger.configure(logger_path, ["tensorboard", "stdout"])
    logger.logtxt(call_params_string_array)

    with open(logger_path + "/params.txt", 'w') as file:
        file.write(json.dumps(call_params_string_array, indent=4))
    copyfile(__file__, logger_path + "/script.txt")

    # Create the training model
    model = deepq.models.mlp(hidden_layers)
    act = deepq.learn(
        env=env,
        q_func=model,
        lr=lr,
        max_timesteps=max_timesteps,
        buffer_size=buffer_size,
        exploration_fraction=exploration_fraction,
        exploration_final_eps=explore_final_eps,
        train_freq=train_freq,
        batch_size=batch_size,
        print_freq=1,
        checkpoint_freq=checkpoint_freq,
        learning_starts=learning_starts,
        gamma=gamma,
        target_network_update_freq=target_network_update_freq,
        prioritized_replay=prioritized_replay,
        prioritized_replay_alpha=prioritized_replay_alpha,
        prioritized_replay_beta0=prioritized_replay_beta0,
        prioritized_replay_beta_iters=prioritized_replay_beta_iters,
        prioritized_replay_eps=prioritized_replay_eps,
        num_cpu=num_cpu,
        param_noise=param_noise,
        callback=None,
        model_path=logger_path
    )
    save_path = logger_path + "/model-" + datetime.now().strftime('%Y-%m-%d_%H-%M-%S') + ".pkl"
    print("Saving last model to {}".format(save_path))
    act.save(save_path)

def main():

    reward_functions = [BaseTraciEnv.reward_average_accumulated_wait_time,
                        BaseTraciEnv.reward_rms_accumulated_wait_time,
                        BaseTraciEnv.reward_total_waiting_vehicles,
                        BaseTraciEnv.reward_total_in_queue_3cross,
                        BaseTraciEnv.reward_arrived_vehicles,
                        BaseTraciEnv.reward_average_speed,
                        BaseTraciEnv.reward_halting_in_queue_3cross]

    for rf in reward_functions:
        print("Now reward function is:", rf)
        g = tf.Graph()
        sess = tf.InteractiveSession(graph=g)
        with g.as_default():
            train_and_log(reward_function=rf)

if __name__ == '__main__':
    main()
<commit_msg><commit_after>import json
import os, sys
from shutil import copyfile

from datetime import datetime
import inspect
import tensorflow as tf

if 'SUMO_HOME' in os.environ:
    tools = os.path.join(os.environ['SUMO_HOME'], 'tools')
    sys.path.append(tools)
else:
    sys.exit("please declare environment variable 'SUMO_HOME'")

import gym

from baselines import deepq
import Traci_1_cross_env.Traci_1_cross_env
import Traci_2_cross_env.Traci_2_cross_env
import Traci_3_cross_env.Traci_3_cross_env
from baselines import logger, logger_utils
from BaseTraciEnv import BaseTraciEnv
from pathlib import Path

def callback(lcl, glb):
    # stop training if reward exceeds 199
    is_solved = lcl['t'] > 100 and sum(lcl['episode_rewards'][-101:-1]) / 100 >= 199
    return is_solved

def train_and_log(environment="Traci_3_cross_env-v0",
                  car_chances=1000,
                  reward_function=BaseTraciEnv.reward_average_accumulated_wait_time,
                  lr=1e-3,
                  max_timesteps=int(1e5),
                  buffer_size=10000,
                  exploration_fraction=0.5,
                  explore_final_eps=0.02,
                  train_freq=100,
                  batch_size=32,
                  checkpoint_freq=int(10000),
                  learning_starts=10000,
                  gamma=0.9,
                  target_network_update_freq=500,
                  start_car_probabilities=[0.25, 0.05],  # [0.1,0.1,0.1,0.1,0.1,0.1,0.1], #For traci_3_cross: Bigroad_spawn_prob,Smallroad_spawn_prob
                  end_car_probabilities=[1.0,0.1],
                  num_steps_from_start_car_probs_to_end_car_probs=1e5,
                  prioritized_replay=False,
                  prioritized_replay_alpha=0.6,
                  prioritized_replay_beta0=0.4,
                  prioritized_replay_beta_iters=None,
                  prioritized_replay_eps=1e-6,
                  num_cpu=4,
                  param_noise=False,
                  state_use_queue_length=True,
                  state_use_tl_state=True,
                  state_use_time_since_tl_change=True,
                  state_use_avg_speed=False,
                  hidden_layers=[8, 8, 8],
                  num_actions_pr_trafficlight=3):
    # Print call values
    frame = inspect.currentframe()
    args, _, _, values = inspect.getargvalues(frame)

    call_params_string_array = ['function name "%s"' % inspect.getframeinfo(frame)[2]]
    for i in args:
        call_params_string_array.append("    %s = %s" % (i, values[i]))

    # Setup path of logging, name of environment and save the current arguments (this script)
    log_dir = [os.path.join(str(Path.home()), "Desktop"), environment]
    logger_path = logger_utils.path_with_date(log_dir[0], log_dir[1])

    # Create environment and initialize
    env = gym.make(log_dir[1])
    env.configure_traci(num_car_chances=car_chances,
                        start_car_probabilities=start_car_probabilities,
                        end_car_probabilities=end_car_probabilities,
                        num_steps_from_start_car_probs_to_end_car_probs=num_steps_from_start_car_probs_to_end_car_probs,
                        reward_func=reward_function,
                        state_contain_num_cars_in_queue_history=state_use_queue_length,
                        state_contain_avg_speed_between_detectors_history=state_use_avg_speed,
                        state_contain_time_since_tl_change=state_use_time_since_tl_change,
                        state_contain_tl_state_history=state_use_tl_state,
                        num_actions_pr_trafficlight=num_actions_pr_trafficlight)
    #env.render()

    # Initialize logger
    logger.reset()
    logger.configure(logger_path, ["tensorboard", "stdout"])
    logger.logtxt(call_params_string_array)

    with open(logger_path + "/params.txt", 'w') as file:
        file.write(json.dumps(call_params_string_array, indent=4))
    copyfile(__file__, logger_path + "/script.txt")

    # Create the training model
    model = deepq.models.mlp(hidden_layers)
    act = deepq.learn(
        env=env,
        q_func=model,
        lr=lr,
        max_timesteps=max_timesteps,
        buffer_size=buffer_size,
        exploration_fraction=exploration_fraction,
        exploration_final_eps=explore_final_eps,
        train_freq=train_freq,
        batch_size=batch_size,
        print_freq=1,
        checkpoint_freq=checkpoint_freq,
        learning_starts=learning_starts,
        gamma=gamma,
        target_network_update_freq=target_network_update_freq,
        prioritized_replay=prioritized_replay,
        prioritized_replay_alpha=prioritized_replay_alpha,
        prioritized_replay_beta0=prioritized_replay_beta0,
        prioritized_replay_beta_iters=prioritized_replay_beta_iters,
        prioritized_replay_eps=prioritized_replay_eps,
        num_cpu=num_cpu,
        param_noise=param_noise,
        callback=None,
        model_path=logger_path
    )
    save_path = logger_path + "/model-" + datetime.now().strftime('%Y-%m-%d_%H-%M-%S') + ".pkl"
    print("Saving last model to {}".format(save_path))
    act.save(save_path)

def main():

    reward_functions = [BaseTraciEnv.reward_average_accumulated_wait_time,
                        BaseTraciEnv.reward_rms_accumulated_wait_time,
                        BaseTraciEnv.reward_total_waiting_vehicles,
                        BaseTraciEnv.reward_total_in_queue_3cross,
                        BaseTraciEnv.reward_arrived_vehicles,
                        BaseTraciEnv.reward_average_speed,
                        BaseTraciEnv.reward_halting_in_queue_3cross]

    for rf in reward_functions:
        print("Now reward function is:", rf)
        g = tf.Graph()
        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True

        sess = tf.InteractiveSession(graph=g,config=config)
        with g.as_default():
            train_and_log(reward_function=rf)

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


<commit_before>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import math

import measurement_stats as mstats
from measurement_stats.angle import Angle

from tracksim.trackway import TrackPosition

class LineSegment2D(object):
    """A class for..."""

    def __init__(self, start: TrackPosition, end: TrackPosition):
        """Creates a new instance of LineSegment2D."""
        self.start = start.clone()
        self.end   = end.clone()

    @property
    def is_value(self):
        return bool(
            abs(self.end.x - self.start.x) + abs(self.end.y - self.start.y)
        )

    @property
    def length(self):
        """ The length of the line segment. """
        return self.start.distance_between(self.end)

    @property
    def angle(self):
        v = self.end.clone()
        v.x -= self.start.x
        v.y -= self.start.y

        vxUnc = (
            self.start.x.raw_uncertainty ** 2 +
            self.end.x.raw_uncertainty ** 2
        ) ** 0.5

        vyUnc = (
            self.start.y.raw_uncertainty ** 2 +
            self.end.y.raw_uncertainty ** 2
        ) ** 0.5

        value = math.atan2(v.y, v.x)
        lengthSqr = v.x.raw ** 2 + v.y.raw ** 2
        unc = abs(v.y / lengthSqr) * vxUnc + abs(v.x / lengthSqr) * vyUnc
        return Angle(radians=value, uncertainty=unc)

    @property
    def slope(self):
        """ Returns the slope of the line as a ValueUncertainty named tuple. """
        s       = self.start
        e       = self.end
        deltaX  = e.x - s.x
        deltaY  = e.y - s.y

        try:
            slope   = deltaY / deltaX
            unc = abs(1.0 / deltaX) * (
                s.y.raw_uncertainty + e.y.raw_uncertainty
            ) + abs(slope / deltaX) * (
                s.x.raw_uncertainty + e.x.raw_uncertainty
            )
            return mstats.ValueUncertainty(slope, unc)
        except Exception:
            return None

    @property
    def midpoint(self):
        """
        Returns the midpoint of the line as a PositionValue2D instance.
        """

        x = 0.5 * (self.start.x + self.end.x)
        y = 0.5 * (self.start.y + self.end.y)

        return TrackPosition(x=x, y=y)

    def serialize(self):
        """

        :return:
        """

        return dict(
            start=self.start.serialize(),
            end=self.end.serialize()
        )

    def from_dict(self, source):
        """

        :param source:
        :return:
        """

        self.start = mstats.ValueUncertainty()
        self.start.from_dict(source['start'])

        self.end = mstats.ValueUncertainty()
        self.end.from_dict(source['end'])

    def add_offset(self, point):
        """add_offset doc..."""

        self.start.x += point.x
        self.start.y += point.y
        self.end.x += point.x
        self.end.y += point.y

    def get_parametric_position(self, value, clamp=True):
        """get_parametric_position doc..."""

        if clamp:
            value = max(0.0, min(value, 1.0))

        x = self.start.x + value * (self.end.x - self.start.x)
        y = self.start.y + value * (self.end.y - self.start.y)

        return TrackPosition(x=x, y=y)

    def adjust_point_along_line(self, point, delta, inPlace=False):
        """adjust_point_along_line doc..."""

        if not inPlace:
            point = point.clone()

        remove = mstats.value2D.Point2D(
            x=self.start.x.clone(),
            y=self.start.y.clone()
        )
        remove.invert()
        remove = TrackPosition(x=remove.x, y=remove.y)

        line = self.clone()
        line.add_offset(remove)

        position = line.get_parametric_position(
            delta / self.length.raw,
            clamp=False
        )
        point.add(position)

        return point

    def rotate(self, angle: float, pivot=None):
        """rotate doc..."""

        self.start.rotate(angle, pivot=pivot)
        self.end.rotate(angle, pivot=pivot)

    def angle_between(self, line):
        """
        Returns an Angle instance that represents the angle between this line
        segment and the one specified in the arguments.
        """

        return self.angle.difference_between(line.angle)

    def angle_between_point(self, position):
        """angle_between_point doc..."""
        a = mstats.value2D.Point2D(
            x=self.end.x - self.start.x,
            y=self.end.y - self.start.y
        )

        b = mstats.value2D.Point2D(
            x=self.end.x - self.start.x,
            y=self.end.y - self.start.y
        )

        return b.angle_between(a)

    def clone(self):
        """
        Returns a new LineSegment2D that is a clone of this instance. The start
        and end points are cloned as well making a completely separate copy with
        no reference dependencies on this instance.
        """

        return LineSegment2D(
            start=self.start.clone(),
            end=self.end.clone()
        )

    def distance_to_point(self, point):
        """
        Calculates the smallest distance between the specified point and  this
        line segment using the standard formulation as described in:

            http://en.wikipedia.org \
                /wiki/Distance_from_a_point_to_a_line \
                #Line_defined_by_two_points
        """

        length = self.length.raw
        if not length:
            raise ValueError(
                'Cannot calculate point distance. Invalid line segment.'
            )

        sx = self.start.x.raw
        sx_unc = self.start.x.raw_uncertainty
        sy = self.start.y.raw
        sy_unc = self.start.y.raw_uncertainty

        ex = self.end.x.raw
        ex_unc = self.end.x.raw_uncertainty
        ey = self.end.y.raw
        ey_unc = self.end.y.raw_uncertainty

        px = point.x.raw
        px_unc = point.x.raw_uncertainty
        py = point.y.raw
        py_unc = point.y.raw_uncertainty

        deltaX = ex - sx
        deltaY = ey - sy

        if deltaX == 0.0:
            # Vertical Line
            distance = abs(sx - px)
        elif deltaY == 0.0:
            # Horizontal line
            distance = abs(sy - py)
        else:
            distance = abs(
                deltaY * px - deltaX * py - sx * ey + ex * sy
            ) / length

        B = deltaY * px - deltaX * py - sx * ey + ex * sy
        AbsB = abs(B)
        D = (deltaX * deltaX + deltaY * deltaY) ** 0.5
        DPrime = 1.0 / (deltaX * deltaX + deltaY * deltaY) ** (3.0 / 2.0)
        bBD = B / (AbsB * D)

        pointXErr = px_unc * abs(deltaY * B / (AbsB * D))
        pointYErr = py_unc * abs(deltaX * B / (AbsB * D))
        startXErr = sx_unc * abs(AbsB * DPrime + bBD * (py - ey))
        startYErr = sy_unc * abs(AbsB * DPrime + bBD * (ex - px))
        endXErr = ex_unc * abs(bBD * (sy - py) - AbsB * DPrime)
        endYErr = ey_unc * abs(bBD * (px - sx) - AbsB * DPrime)
        error = (
            pointXErr + pointYErr +
            startXErr + startYErr +
            endXErr + endYErr
        )

        return mstats.ValueUncertainty(distance, error)

    def closest_point_on_line(self, point, contained=True):
        """
        Finds the closest point on a line to the specified point using the
        formulae discussed in the "another formula" section of:
            http://en.m.wikipedia.org \
                /wiki/Distance_from_a_point_to_a_line \
                #Another_formula
        """

        length = self.length.raw
        if not length:
            raise ValueError('Cannot calculate point. Invalid line segment.')

        sx = self.start.x.raw
        sx_unc = self.start.x.raw_uncertainty
        sy = self.start.y.raw
        sy_unc = self.start.y.raw_uncertainty

        ex = self.end.x.raw
        ex_unc = self.end.x.raw_uncertainty
        ey = self.end.y.raw
        ey_unc = self.end.y.raw_uncertainty

        px = point.x.raw
        px_unc = point.x.raw_uncertainty
        py = point.y.raw
        py_unc = point.y.raw_uncertainty

        deltaX = ex - sx
        deltaY = ey - sy
        rotate = False
        slope = 0.0
        slopeUnc = 0.0

        try:
            slope = deltaY / deltaX
            slopeUnc = (
                abs(1.0 / deltaX) * (sy_unc + ey_unc) +
                abs(slope / deltaX) * (sx_unc + ex_unc)
            )
        except Exception:
            rotate = True

        if rotate or (abs(slope) > 1.0 and abs(slopeUnc / slope) > 0.5):
            a = Angle(degrees=20.0)
            line = self.clone()
            line.rotate(a.value, self.start)
            p = point.clone()
            p.rotate(a, self.start)
            result = line.closest_point_on_line(p, contained=contained)
            if result is None:
                return result

            a.degrees = -20.0
            result.rotate(a.radians, self.start)
            return result

        intercept = sy - slope * sx
        denominator = slope * slope + 1.0
        numerator = point.x + slope * (point.y - intercept)

        x = (numerator / denominator).raw
        y = ((slope * numerator) / denominator + intercept).raw

        if contained:
            # Check to see if point is between start and end values
            xRange = sorted([sx, ex])
            yRange = sorted([sy, ey])
            eps = 1e-8
            xMin = x - eps
            xMax = x + eps
            yMin = y - eps
            yMax = y + eps
            outside = (
                xRange[1] < xMin or
                xMax < xRange[0] or
                yRange[1] < yMin or
                yMax < yRange[0]
            )
            if outside:
                return None

        startDist = self.start.distance_from(x, y)
        endDist = self.end.distance_from(x, y)

        x_unc = (
            startDist / length * sx_unc +
            endDist / length * ex_unc
        )
        x_unc = (x_unc ** 2 + px_unc ** 2) ** 0.5

        y_unc = (
            startDist / length * sy_unc +
            endDist / length * ey_unc
        )
        y_unc = (y_unc ** 2 + py_unc ** 2) ** 0.5

        return TrackPosition(
            x=mstats.ValueUncertainty(x, x_unc),
            y=mstats.ValueUncertainty(y, y_unc)
        )

    def post_extend_line(self, lengthAdjust, replace=True):
        """extendLine doc..."""
        newX, newY = self._extrapolate_by_length(lengthAdjust, pre=False)
        if not replace:
            self.end = self.end.clone()
        self.end.x = newX
        self.end.y = newY

    def pre_extend_line(self, lengthAdjust, replace=True):
        """pre_extend_line doc..."""
        newX, newY = self._extrapolate_by_length(lengthAdjust, pre=True)
        if not replace:
            self.start = self.start.clone()
        self.start.x = newX
        self.start.y = newY

    def create_next_line_segment(self, length=None):
        """
        Creates a line segment using this line segment as a guide that starts
        where this segment ends and has the same slope and orientation. The new
        line segment will be of the specified length, or if no length is
        specified the same length as this line segment.
        """

        if length is None:
            length = self.length

        target = self.clone()
        target.post_extend_line(lengthAdjust=length)
        target.start = self.end.clone()

        return target

    def create_previous_line_segment(self, length=None):
        """
        Creates a line segment using this line segment as a guide that ends
        where this segment begins and has the same slope and orientation. The
        new line segment will be of the specified length, or if no length is
        specified the same length as this line segment.
        """

        if length is None:
            length = self.length

        target = self.clone()
        target.pre_extend_line(lengthAdjust=length)
        target.end = self.start.clone()

        return target

    def _extrapolate_by_length(self, lengthAdjust, pre=False):
        """
        _extrapolate_by_length doc...
        """

        sx = self.start.x
        sy = self.start.y

        ex = self.end.x
        ey = self.end.y

        if pre:
            startY = sy
            startX = sx
            point = self.end
        else:
            startY = ey
            startX = ex
            point = self.start

        deltaX = startX - point.x
        deltaY = startY - point.y

        try:
            if mstats.value.equivalent(deltaX.value, 0.0):
                direction = deltaY.raw / abs(deltaY.raw)
                return startX, startY + direction * lengthAdjust

            if mstats.value.equivalent(deltaY.value, 0.0):
                direction = deltaX.raw / abs(deltaX.raw)
                return startX + direction * lengthAdjust, startY
        except ZeroDivisionError as err:
            print('Zero Division Error:')
            print('Delta X: {} Delta Y: {}'.format(deltaX, deltaY))
            print('Line Start: ({}, {})'.format(self.start.x, self.start.y))
            print('Line End: ({}, {})'.format(self.end.x, self.end.y))
            raise

        deltaX = startX - point.x
        deltaY = startY - point.y

        angle = math.atan2(deltaY.raw, deltaX.raw)

        return (
            startX + lengthAdjust * math.cos(angle),
            startY + lengthAdjust * math.sin(angle)
        )
<commit_msg>Fixes another angle argument passing in geometry.py<commit_after>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import math

import measurement_stats as mstats
from measurement_stats.angle import Angle

from tracksim.trackway import TrackPosition

class LineSegment2D(object):
    """A class for..."""

    def __init__(self, start: TrackPosition, end: TrackPosition):
        """Creates a new instance of LineSegment2D."""
        self.start = start.clone()
        self.end   = end.clone()

    @property
    def is_value(self):
        return bool(
            abs(self.end.x - self.start.x) + abs(self.end.y - self.start.y)
        )

    @property
    def length(self):
        """ The length of the line segment. """
        return self.start.distance_between(self.end)

    @property
    def angle(self):
        v = self.end.clone()
        v.x -= self.start.x
        v.y -= self.start.y

        vxUnc = (
            self.start.x.raw_uncertainty ** 2 +
            self.end.x.raw_uncertainty ** 2
        ) ** 0.5

        vyUnc = (
            self.start.y.raw_uncertainty ** 2 +
            self.end.y.raw_uncertainty ** 2
        ) ** 0.5

        value = math.atan2(v.y, v.x)
        lengthSqr = v.x.raw ** 2 + v.y.raw ** 2
        unc = abs(v.y / lengthSqr) * vxUnc + abs(v.x / lengthSqr) * vyUnc
        return Angle(radians=value, uncertainty=unc)

    @property
    def slope(self):
        """ Returns the slope of the line as a ValueUncertainty named tuple. """
        s       = self.start
        e       = self.end
        deltaX  = e.x - s.x
        deltaY  = e.y - s.y

        try:
            slope   = deltaY / deltaX
            unc = abs(1.0 / deltaX) * (
                s.y.raw_uncertainty + e.y.raw_uncertainty
            ) + abs(slope / deltaX) * (
                s.x.raw_uncertainty + e.x.raw_uncertainty
            )
            return mstats.ValueUncertainty(slope, unc)
        except Exception:
            return None

    @property
    def midpoint(self):
        """
        Returns the midpoint of the line as a PositionValue2D instance.
        """

        x = 0.5 * (self.start.x + self.end.x)
        y = 0.5 * (self.start.y + self.end.y)

        return TrackPosition(x=x, y=y)

    def serialize(self):
        """

        :return:
        """

        return dict(
            start=self.start.serialize(),
            end=self.end.serialize()
        )

    def from_dict(self, source):
        """

        :param source:
        :return:
        """

        self.start = mstats.ValueUncertainty()
        self.start.from_dict(source['start'])

        self.end = mstats.ValueUncertainty()
        self.end.from_dict(source['end'])

    def add_offset(self, point):
        """add_offset doc..."""

        self.start.x += point.x
        self.start.y += point.y
        self.end.x += point.x
        self.end.y += point.y

    def get_parametric_position(self, value, clamp=True):
        """get_parametric_position doc..."""

        if clamp:
            value = max(0.0, min(value, 1.0))

        x = self.start.x + value * (self.end.x - self.start.x)
        y = self.start.y + value * (self.end.y - self.start.y)

        return TrackPosition(x=x, y=y)

    def adjust_point_along_line(self, point, delta, inPlace=False):
        """adjust_point_along_line doc..."""

        if not inPlace:
            point = point.clone()

        remove = mstats.value2D.Point2D(
            x=self.start.x.clone(),
            y=self.start.y.clone()
        )
        remove.invert()
        remove = TrackPosition(x=remove.x, y=remove.y)

        line = self.clone()
        line.add_offset(remove)

        position = line.get_parametric_position(
            delta / self.length.raw,
            clamp=False
        )
        point.add(position)

        return point

    def rotate(self, angle: float, pivot=None):
        """rotate doc..."""

        self.start.rotate(angle, pivot=pivot)
        self.end.rotate(angle, pivot=pivot)

    def angle_between(self, line):
        """
        Returns an Angle instance that represents the angle between this line
        segment and the one specified in the arguments.
        """

        return self.angle.difference_between(line.angle)

    def angle_between_point(self, position):
        """angle_between_point doc..."""
        a = mstats.value2D.Point2D(
            x=self.end.x - self.start.x,
            y=self.end.y - self.start.y
        )

        b = mstats.value2D.Point2D(
            x=self.end.x - self.start.x,
            y=self.end.y - self.start.y
        )

        return b.angle_between(a)

    def clone(self):
        """
        Returns a new LineSegment2D that is a clone of this instance. The start
        and end points are cloned as well making a completely separate copy with
        no reference dependencies on this instance.
        """

        return LineSegment2D(
            start=self.start.clone(),
            end=self.end.clone()
        )

    def distance_to_point(self, point):
        """
        Calculates the smallest distance between the specified point and  this
        line segment using the standard formulation as described in:

            http://en.wikipedia.org \
                /wiki/Distance_from_a_point_to_a_line \
                #Line_defined_by_two_points
        """

        length = self.length.raw
        if not length:
            raise ValueError(
                'Cannot calculate point distance. Invalid line segment.'
            )

        sx = self.start.x.raw
        sx_unc = self.start.x.raw_uncertainty
        sy = self.start.y.raw
        sy_unc = self.start.y.raw_uncertainty

        ex = self.end.x.raw
        ex_unc = self.end.x.raw_uncertainty
        ey = self.end.y.raw
        ey_unc = self.end.y.raw_uncertainty

        px = point.x.raw
        px_unc = point.x.raw_uncertainty
        py = point.y.raw
        py_unc = point.y.raw_uncertainty

        deltaX = ex - sx
        deltaY = ey - sy

        if deltaX == 0.0:
            # Vertical Line
            distance = abs(sx - px)
        elif deltaY == 0.0:
            # Horizontal line
            distance = abs(sy - py)
        else:
            distance = abs(
                deltaY * px - deltaX * py - sx * ey + ex * sy
            ) / length

        B = deltaY * px - deltaX * py - sx * ey + ex * sy
        AbsB = abs(B)
        D = (deltaX * deltaX + deltaY * deltaY) ** 0.5
        DPrime = 1.0 / (deltaX * deltaX + deltaY * deltaY) ** (3.0 / 2.0)
        bBD = B / (AbsB * D)

        pointXErr = px_unc * abs(deltaY * B / (AbsB * D))
        pointYErr = py_unc * abs(deltaX * B / (AbsB * D))
        startXErr = sx_unc * abs(AbsB * DPrime + bBD * (py - ey))
        startYErr = sy_unc * abs(AbsB * DPrime + bBD * (ex - px))
        endXErr = ex_unc * abs(bBD * (sy - py) - AbsB * DPrime)
        endYErr = ey_unc * abs(bBD * (px - sx) - AbsB * DPrime)
        error = (
            pointXErr + pointYErr +
            startXErr + startYErr +
            endXErr + endYErr
        )

        return mstats.ValueUncertainty(distance, error)

    def closest_point_on_line(self, point, contained=True):
        """
        Finds the closest point on a line to the specified point using the
        formulae discussed in the "another formula" section of:
            http://en.m.wikipedia.org \
                /wiki/Distance_from_a_point_to_a_line \
                #Another_formula
        """

        length = self.length.raw
        if not length:
            raise ValueError('Cannot calculate point. Invalid line segment.')

        sx = self.start.x.raw
        sx_unc = self.start.x.raw_uncertainty
        sy = self.start.y.raw
        sy_unc = self.start.y.raw_uncertainty

        ex = self.end.x.raw
        ex_unc = self.end.x.raw_uncertainty
        ey = self.end.y.raw
        ey_unc = self.end.y.raw_uncertainty

        px = point.x.raw
        px_unc = point.x.raw_uncertainty
        py = point.y.raw
        py_unc = point.y.raw_uncertainty

        deltaX = ex - sx
        deltaY = ey - sy
        rotate = False
        slope = 0.0
        slopeUnc = 0.0

        try:
            slope = deltaY / deltaX
            slopeUnc = (
                abs(1.0 / deltaX) * (sy_unc + ey_unc) +
                abs(slope / deltaX) * (sx_unc + ex_unc)
            )
        except Exception:
            rotate = True

        if rotate or (abs(slope) > 1.0 and abs(slopeUnc / slope) > 0.5):
            a = Angle(degrees=20.0)
            line = self.clone()
            line.rotate(a.radians, self.start)
            p = point.clone()
            p.rotate(a.radians, self.start)
            result = line.closest_point_on_line(p, contained=contained)
            if result is None:
                return result

            a.degrees = -20.0
            result.rotate(a.radians, self.start)
            return result

        intercept = sy - slope * sx
        denominator = slope * slope + 1.0
        numerator = point.x + slope * (point.y - intercept)

        x = (numerator / denominator).raw
        y = ((slope * numerator) / denominator + intercept).raw

        if contained:
            # Check to see if point is between start and end values
            xRange = sorted([sx, ex])
            yRange = sorted([sy, ey])
            eps = 1e-8
            xMin = x - eps
            xMax = x + eps
            yMin = y - eps
            yMax = y + eps
            outside = (
                xRange[1] < xMin or
                xMax < xRange[0] or
                yRange[1] < yMin or
                yMax < yRange[0]
            )
            if outside:
                return None

        startDist = self.start.distance_from(x, y)
        endDist = self.end.distance_from(x, y)

        x_unc = (
            startDist / length * sx_unc +
            endDist / length * ex_unc
        )
        x_unc = (x_unc ** 2 + px_unc ** 2) ** 0.5

        y_unc = (
            startDist / length * sy_unc +
            endDist / length * ey_unc
        )
        y_unc = (y_unc ** 2 + py_unc ** 2) ** 0.5

        return TrackPosition(
            x=mstats.ValueUncertainty(x, x_unc),
            y=mstats.ValueUncertainty(y, y_unc)
        )

    def post_extend_line(self, lengthAdjust, replace=True):
        """extendLine doc..."""
        newX, newY = self._extrapolate_by_length(lengthAdjust, pre=False)
        if not replace:
            self.end = self.end.clone()
        self.end.x = newX
        self.end.y = newY

    def pre_extend_line(self, lengthAdjust, replace=True):
        """pre_extend_line doc..."""
        newX, newY = self._extrapolate_by_length(lengthAdjust, pre=True)
        if not replace:
            self.start = self.start.clone()
        self.start.x = newX
        self.start.y = newY

    def create_next_line_segment(self, length=None):
        """
        Creates a line segment using this line segment as a guide that starts
        where this segment ends and has the same slope and orientation. The new
        line segment will be of the specified length, or if no length is
        specified the same length as this line segment.
        """

        if length is None:
            length = self.length

        target = self.clone()
        target.post_extend_line(lengthAdjust=length)
        target.start = self.end.clone()

        return target

    def create_previous_line_segment(self, length=None):
        """
        Creates a line segment using this line segment as a guide that ends
        where this segment begins and has the same slope and orientation. The
        new line segment will be of the specified length, or if no length is
        specified the same length as this line segment.
        """

        if length is None:
            length = self.length

        target = self.clone()
        target.pre_extend_line(lengthAdjust=length)
        target.end = self.start.clone()

        return target

    def _extrapolate_by_length(self, lengthAdjust, pre=False):
        """
        _extrapolate_by_length doc...
        """

        sx = self.start.x
        sy = self.start.y

        ex = self.end.x
        ey = self.end.y

        if pre:
            startY = sy
            startX = sx
            point = self.end
        else:
            startY = ey
            startX = ex
            point = self.start

        deltaX = startX - point.x
        deltaY = startY - point.y

        try:
            if mstats.value.equivalent(deltaX.value, 0.0):
                direction = deltaY.raw / abs(deltaY.raw)
                return startX, startY + direction * lengthAdjust

            if mstats.value.equivalent(deltaY.value, 0.0):
                direction = deltaX.raw / abs(deltaX.raw)
                return startX + direction * lengthAdjust, startY
        except ZeroDivisionError as err:
            print('Zero Division Error:')
            print('Delta X: {} Delta Y: {}'.format(deltaX, deltaY))
            print('Line Start: ({}, {})'.format(self.start.x, self.start.y))
            print('Line End: ({}, {})'.format(self.end.x, self.end.y))
            raise

        deltaX = startX - point.x
        deltaY = startY - point.y

        angle = math.atan2(deltaY.raw, deltaX.raw)

        return (
            startX + lengthAdjust * math.cos(angle),
            startY + lengthAdjust * math.sin(angle)
        )
<|endoftext|>


<commit_before>from __future__ import print_function, division
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider
from ciabatta.ejm_rcparams import reds_cmap
from agaro import output_utils
from ahoy.utils import utils
from ahoy.plot.var_plot import VarPlot
import ahoy.fields

def plot_2d(dirname):
    fig = plt.figure()
    ax_vis = fig.add_subplot(111)

    fnames = output_utils.get_filenames(dirname)
    m_0 = output_utils.filename_to_model(fnames[0])

    L = m_0.ships.agents.positions.L

    ax_vis.set_xlim(-L[0] / 2.0, L[0] / 2.0)
    ax_vis.set_ylim(-L[1] / 2.0, L[1] / 2.0)
    ax_vis.set_aspect('equal')

    plt.subplots_adjust(left=0.25, bottom=0.25)
    has_c_field = isinstance(m_0.ships.c_field, ahoy.fields.FoodField)
    if has_c_field:
        plot_c = VarPlot(m_0.ships.c_field.c, cmap=reds_cmap, axes=ax_vis)
    plot_p = ax_vis.quiver(m_0.ships.agents.positions.r_w[:, 0],
                           m_0.ships.agents.positions.r_w[:, 1],
                           m_0.ships.agents.directions.u[:, 0],
                           m_0.ships.agents.directions.u[:, 1])

    ax_slide = plt.axes([0.25, 0.1, 0.65, 0.03])
    t_slider = Slider(ax_slide, 'Time', 0, len(fnames), valinit=0)

    t_time = fig.text(0.1, 0.5, '')

    def update(val):
        fname_i = int(round(val))
        if 0 <= fname_i < len(fnames):
            m = output_utils.filename_to_model(fnames[fname_i])
            if has_c_field:
                plot_c.update(m.ships.c_field.c)
            plot_p.set_offsets(m.ships.agents.positions.r_w)
            plot_p.set_UVC(m.ships.agents.directions.u[:, 0],
                           m.ships.agents.directions.u[:, 1])
            t_time.set_text('Time: {:g}'.format(m.ships.time.t))

            fig.canvas.draw_idle()

    t_slider.on_changed(update)

    plt.show()

def plot_linear_density(dirname):
    fig = plt.figure()
    ax_d = fig.add_subplot(211)
    ax_c = fig.add_subplot(212)

    fnames = output_utils.get_filenames(dirname)
    m_0 = output_utils.filename_to_model(fnames[0])

    L = m_0.ships.agents.positions.L

    dx = L[0] / 100.0

    plt.subplots_adjust(left=0.25, bottom=0.25)

    ds, xbs = utils.get_linear_density(m_0, dx)

    plot_d = ax_d.bar(xbs[:-1], ds, width=xbs[1] - xbs[0])
    c_field = m_0.ships.c_field.c
    plot_c = ax_c.scatter(c_field.mesh.cellCenters[0, :], c_field.value)
    ax_slide = plt.axes([0.25, 0.1, 0.65, 0.03])
    t_slider = Slider(ax_slide, 'Index', 0, len(fnames), valinit=0)

    ax_d.set_xlim(-L[0] / 2.0, L[0] / 2.0)
    ax_c.set_xlim(-L[0] / 2.0, L[0] / 2.0)
    ax_c.set_ylim(0.0, m_0.ships.c_field.c_0)

    def update(val):
        fname_i = int(round(val))
        if 0 <= fname_i < len(fnames):
            m = output_utils.filename_to_model(fnames[fname_i])
            ds, xbs = utils.get_linear_density(m, dx)
            for rect, d in zip(plot_d, ds):
                rect.set_height(d)
            c_field = m.ships.c_field.c
            plot_c.set_offsets(np.array([c_field.mesh.cellCenters[0, :],
                                         c_field.value]).T)
            fig.canvas.draw_idle()

    t_slider.on_changed(update)

    plt.show()

def plot_1d(dirname):
    fig = plt.figure()
    ax_vis = fig.add_subplot(211)
    ax_d = fig.add_subplot(212)

    fnames = output_utils.get_filenames(dirname)
    m_0 = output_utils.filename_to_model(fnames[0])

    L = m_0.ships.agents.positions.L

    ax_vis.set_xlim(-L[0] / 2.0, L[0] / 2.0)
    ax_d.set_xlim(-L[0] / 2.0, L[0] / 2.0)

    dx = L / 100.0

    plt.subplots_adjust(left=0.25, bottom=0.25)
    plot_p = ax_vis.scatter(m_0.ships.agents.positions.r_w[:, 0],
                            np.zeros([m_0.ships.agents.n]))

    d = m_0.ships.agents.positions.get_density_field(dx)
    x = np.linspace(-L[0] / 2.0, L[0] / 2.0, d.shape[0])

    plot_d = ax_d.bar(x, d, width=x[1] - x[0])

    ax_slide = plt.axes([0.25, 0.1, 0.65, 0.03])
    t_slider = Slider(ax_slide, 'Index', 0, len(fnames), valinit=0)

    def update(val):
        fname_i = int(round(val))
        if 0 <= fname_i < len(fnames):
            m = output_utils.filename_to_model(fnames[fname_i])
            plot_p.set_offsets(np.array([m.ships.agents.positions.r_w[:, 0],
                                         np.zeros([m.ships.agents.n])]).T)
            ds = m.ships.agents.positions.get_density_field(dx)
            for rect, d in zip(plot_d, ds):
                rect.set_height(d)
            ax_d.set_ylim(0.0, 1.05 * ds.max())
            fig.canvas.draw_idle()

    t_slider.on_changed(update)

    plt.show()

def plot_vis(dirname):
    dim = output_utils.get_recent_model(dirname).ships.dim
    if dim == 1:
        plot_1d(dirname)
    elif dim == 2:
        plot_2d(dirname)

def plot_t_uds_scalar(dirname, ax):
    ts, uds = utils.t_uds_scalar(dirname)
    ax.plot(ts, uds)

def plot_t_uds_vector(dirname, ax):
    ts, uds = utils.t_uds_vector(dirname)
    for vd_set in uds.T:
        ax.plot(ts, vd_set)

def plot_t_Ds_scalar(dirname, ax):
    ts, Ds = utils.t_Ds_scalar(dirname)
    ax.plot(ts, Ds)

def plot_t_Ds_vector(dirname, ax):
    ts, Ds = utils.t_Ds_vector(dirname)
    for D_set in Ds.T:
        ax.plot(ts, D_set)

def plot_t_rs_scalar(dirname, ax):
    ts, rs = utils.t_rs_scalar(dirname)
    ax.plot(ts, rs)

def plot_t_rs_vector(dirname, ax):
    ts, rs = utils.t_rs_vector(dirname)
    for r_set in rs.T:
        ax.plot(ts, r_set)

def plot_t_u_nets_scalar(dirname, ax):
    ts, u_nets = utils.t_u_nets_scalar(dirname)
    ax.plot(ts, u_nets)

def plot_t_u_nets_vector(dirname, ax):
    ts, u_nets = utils.t_u_nets_vector(dirname)
    for u_net_set in u_nets.T:
        ax.plot(ts, u_net_set)

def plot_chi_uds_x(dirnames, ax):
    chis, uds = utils.chi_uds_x(dirnames)
    i_sort = np.argsort(chis)
    chis, uds = chis[i_sort], uds[i_sort]
    ax.scatter(chis, uds)
    ax.set_xlim(-0.02, 1.0)
    ax.set_ylim(0.0, 1.1)

def plot_pf_Ds_scalar(dirnames, ax):
    pfs, Ds = utils.pf_Ds_scalar(dirnames)
    i_sort = np.argsort(pfs)
    pfs, Ds = pfs[i_sort], Ds[i_sort]
    ax.scatter(pfs, Ds)
    ax.set_xlim(-0.02, 1.0)
    ax.set_ylim(0.0, 410.0)

def plot_pf_uds_x(dirnames, ax):
    pfs, uds = utils.pf_uds_x(dirnames)
    i_sort = np.argsort(pfs)
    pfs, uds = pfs[i_sort], uds[i_sort]
    ax.scatter(pfs, uds / uds[0])
    ax.set_xlim(-0.02, 1.0)
    ax.set_ylim(0.0, 1.01)

def plot_pf_duds_x(dirnames, ax):
    pfs, uds = utils.pf_uds_x(dirnames)
    i_sort = np.argsort(pfs)
    pfs, uds = pfs[i_sort], uds[i_sort]
    uds_norm = uds / uds[0]
    duds = np.diff(uds_norm) / pfs[:-1]
    ax.scatter(pfs[:-1], duds)
    ax.set_xlim(-0.02, 1.0)
    # ax.set_ylim(0.0, 1.0)

def plot_Dr_0_Ds_scalar(dirnames, ax):
    Dr_0s, Ds = utils.Dr_0_Ds_scalar(dirnames)
    i_sort = np.argsort(Dr_0s)
    Dr_0s, Ds = Dr_0s[i_sort], Ds[i_sort]
    ax.scatter(Dr_0s, Ds)
    # ax.set_xlim(-0.02, 1.0)
    # ax.set_ylim(0.0, 410.0)
    ax.set_xscale('log')
    ax.set_yscale('log')

def plot_p_0_Ds_scalar(dirnames, ax):
    p_0s, Ds = utils.p_0_Ds_scalar(dirnames)
    i_sort = np.argsort(p_0s)
    p_0s, Ds = p_0s[i_sort], Ds[i_sort]
    ax.scatter(p_0s, Ds)
    # ax.set_xlim(-0.02, 1.0)
    # ax.set_ylim(0.0, 410.0)
    ax.set_xscale('log')
    ax.set_yscale('log')

def plot_th_density(ds, th_bins, fig):
    ax = plt.subplot(111, polar=True)
    dth = th_bins[1] - th_bins[0]
    ax.plot(th_bins[:-1], ds * dth)
    ax.set_rmax(1.0)
    ax.grid(True)
<commit_msg>Remove unused, specific axis limits in plot functions<commit_after>from __future__ import print_function, division
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider
from ciabatta.ejm_rcparams import reds_cmap
from agaro import output_utils
from ahoy.utils import utils
from ahoy.plot.var_plot import VarPlot
import ahoy.fields

def plot_2d(dirname):
    fig = plt.figure()
    ax_vis = fig.add_subplot(111)

    fnames = output_utils.get_filenames(dirname)
    m_0 = output_utils.filename_to_model(fnames[0])

    L = m_0.ships.agents.positions.L

    ax_vis.set_xlim(-L[0] / 2.0, L[0] / 2.0)
    ax_vis.set_ylim(-L[1] / 2.0, L[1] / 2.0)
    ax_vis.set_aspect('equal')

    plt.subplots_adjust(left=0.25, bottom=0.25)
    has_c_field = isinstance(m_0.ships.c_field, ahoy.fields.FoodField)
    if has_c_field:
        plot_c = VarPlot(m_0.ships.c_field.c, cmap=reds_cmap, axes=ax_vis)
    plot_p = ax_vis.quiver(m_0.ships.agents.positions.r_w[:, 0],
                           m_0.ships.agents.positions.r_w[:, 1],
                           m_0.ships.agents.directions.u[:, 0],
                           m_0.ships.agents.directions.u[:, 1])

    ax_slide = plt.axes([0.25, 0.1, 0.65, 0.03])
    t_slider = Slider(ax_slide, 'Time', 0, len(fnames), valinit=0)

    t_time = fig.text(0.1, 0.5, '')

    def update(val):
        fname_i = int(round(val))
        if 0 <= fname_i < len(fnames):
            m = output_utils.filename_to_model(fnames[fname_i])
            if has_c_field:
                plot_c.update(m.ships.c_field.c)
            plot_p.set_offsets(m.ships.agents.positions.r_w)
            plot_p.set_UVC(m.ships.agents.directions.u[:, 0],
                           m.ships.agents.directions.u[:, 1])
            t_time.set_text('Time: {:g}'.format(m.ships.time.t))

            fig.canvas.draw_idle()

    t_slider.on_changed(update)

    plt.show()

def plot_linear_density(dirname):
    fig = plt.figure()
    ax_d = fig.add_subplot(211)
    ax_c = fig.add_subplot(212)

    fnames = output_utils.get_filenames(dirname)
    m_0 = output_utils.filename_to_model(fnames[0])

    L = m_0.ships.agents.positions.L

    dx = L[0] / 100.0

    plt.subplots_adjust(left=0.25, bottom=0.25)

    ds, xbs = utils.get_linear_density(m_0, dx)

    plot_d = ax_d.bar(xbs[:-1], ds, width=xbs[1] - xbs[0])
    c_field = m_0.ships.c_field.c
    plot_c = ax_c.scatter(c_field.mesh.cellCenters[0, :], c_field.value)
    ax_slide = plt.axes([0.25, 0.1, 0.65, 0.03])
    t_slider = Slider(ax_slide, 'Index', 0, len(fnames), valinit=0)

    ax_d.set_xlim(-L[0] / 2.0, L[0] / 2.0)
    ax_c.set_xlim(-L[0] / 2.0, L[0] / 2.0)
    ax_c.set_ylim(0.0, m_0.ships.c_field.c_0)

    def update(val):
        fname_i = int(round(val))
        if 0 <= fname_i < len(fnames):
            m = output_utils.filename_to_model(fnames[fname_i])
            ds, xbs = utils.get_linear_density(m, dx)
            for rect, d in zip(plot_d, ds):
                rect.set_height(d)
            c_field = m.ships.c_field.c
            plot_c.set_offsets(np.array([c_field.mesh.cellCenters[0, :],
                                         c_field.value]).T)
            fig.canvas.draw_idle()

    t_slider.on_changed(update)

    plt.show()

def plot_1d(dirname):
    fig = plt.figure()
    ax_vis = fig.add_subplot(211)
    ax_d = fig.add_subplot(212)

    fnames = output_utils.get_filenames(dirname)
    m_0 = output_utils.filename_to_model(fnames[0])

    L = m_0.ships.agents.positions.L

    ax_vis.set_xlim(-L[0] / 2.0, L[0] / 2.0)
    ax_d.set_xlim(-L[0] / 2.0, L[0] / 2.0)

    dx = L / 100.0

    plt.subplots_adjust(left=0.25, bottom=0.25)
    plot_p = ax_vis.scatter(m_0.ships.agents.positions.r_w[:, 0],
                            np.zeros([m_0.ships.agents.n]))

    d = m_0.ships.agents.positions.get_density_field(dx)
    x = np.linspace(-L[0] / 2.0, L[0] / 2.0, d.shape[0])

    plot_d = ax_d.bar(x, d, width=x[1] - x[0])

    ax_slide = plt.axes([0.25, 0.1, 0.65, 0.03])
    t_slider = Slider(ax_slide, 'Index', 0, len(fnames), valinit=0)

    def update(val):
        fname_i = int(round(val))
        if 0 <= fname_i < len(fnames):
            m = output_utils.filename_to_model(fnames[fname_i])
            plot_p.set_offsets(np.array([m.ships.agents.positions.r_w[:, 0],
                                         np.zeros([m.ships.agents.n])]).T)
            ds = m.ships.agents.positions.get_density_field(dx)
            for rect, d in zip(plot_d, ds):
                rect.set_height(d)
            ax_d.set_ylim(0.0, 1.05 * ds.max())
            fig.canvas.draw_idle()

    t_slider.on_changed(update)

    plt.show()

def plot_vis(dirname):
    dim = output_utils.get_recent_model(dirname).ships.dim
    if dim == 1:
        plot_1d(dirname)
    elif dim == 2:
        plot_2d(dirname)

def plot_t_uds_scalar(dirname, ax):
    ts, uds = utils.t_uds_scalar(dirname)
    ax.plot(ts, uds)

def plot_t_uds_vector(dirname, ax):
    ts, uds = utils.t_uds_vector(dirname)
    for vd_set in uds.T:
        ax.plot(ts, vd_set)

def plot_t_Ds_scalar(dirname, ax):
    ts, Ds = utils.t_Ds_scalar(dirname)
    ax.plot(ts, Ds)

def plot_t_Ds_vector(dirname, ax):
    ts, Ds = utils.t_Ds_vector(dirname)
    for D_set in Ds.T:
        ax.plot(ts, D_set)

def plot_t_rs_scalar(dirname, ax):
    ts, rs = utils.t_rs_scalar(dirname)
    ax.plot(ts, rs)

def plot_t_rs_vector(dirname, ax):
    ts, rs = utils.t_rs_vector(dirname)
    for r_set in rs.T:
        ax.plot(ts, r_set)

def plot_t_u_nets_scalar(dirname, ax):
    ts, u_nets = utils.t_u_nets_scalar(dirname)
    ax.plot(ts, u_nets)

def plot_t_u_nets_vector(dirname, ax):
    ts, u_nets = utils.t_u_nets_vector(dirname)
    for u_net_set in u_nets.T:
        ax.plot(ts, u_net_set)

def plot_chi_uds_x(dirnames, ax):
    chis, uds = utils.chi_uds_x(dirnames)
    i_sort = np.argsort(chis)
    chis, uds = chis[i_sort], uds[i_sort]
    ax.scatter(chis, uds)
    ax.set_xlim(-0.02, 1.0)
    ax.set_ylim(0.0, 1.1)

def plot_pf_Ds_scalar(dirnames, ax):
    pfs, Ds = utils.pf_Ds_scalar(dirnames)
    i_sort = np.argsort(pfs)
    pfs, Ds = pfs[i_sort], Ds[i_sort]
    ax.scatter(pfs, Ds)
    ax.set_xlim(-0.02, 1.0)

def plot_pf_uds_x(dirnames, ax):
    pfs, uds = utils.pf_uds_x(dirnames)
    i_sort = np.argsort(pfs)
    pfs, uds = pfs[i_sort], uds[i_sort]
    ax.scatter(pfs, uds / uds[0])
    ax.set_xlim(-0.02, 1.0)
    ax.set_ylim(0.0, 1.01)

def plot_pf_duds_x(dirnames, ax):
    pfs, uds = utils.pf_uds_x(dirnames)
    i_sort = np.argsort(pfs)
    pfs, uds = pfs[i_sort], uds[i_sort]
    uds_norm = uds / uds[0]
    duds = np.diff(uds_norm) / pfs[:-1]
    ax.scatter(pfs[:-1], duds)
    ax.set_xlim(-0.02, 1.0)
    # ax.set_ylim(0.0, 1.0)

def plot_Dr_0_Ds_scalar(dirnames, ax):
    Dr_0s, Ds = utils.Dr_0_Ds_scalar(dirnames)
    i_sort = np.argsort(Dr_0s)
    Dr_0s, Ds = Dr_0s[i_sort], Ds[i_sort]
    ax.scatter(Dr_0s, Ds)
    ax.set_xscale('log')
    ax.set_yscale('log')

def plot_p_0_Ds_scalar(dirnames, ax):
    p_0s, Ds = utils.p_0_Ds_scalar(dirnames)
    i_sort = np.argsort(p_0s)
    p_0s, Ds = p_0s[i_sort], Ds[i_sort]
    ax.scatter(p_0s, Ds)
    ax.set_xscale('log')
    ax.set_yscale('log')

def plot_th_density(ds, th_bins, fig):
    ax = plt.subplot(111, polar=True)
    dth = th_bins[1] - th_bins[0]
    ax.plot(th_bins[:-1], ds * dth)
    ax.set_rmax(1.0)
    ax.grid(True)
<|endoftext|>


<commit_before># Copyright 2017 John Reese
# Licensed under the MIT license

import asyncio
import logging
import sqlite3

from concurrent.futures import ThreadPoolExecutor
from functools import partial
from typing import Any, Callable, Iterable

__all__ = [
    'connect',
    'Connection',
    'Cursor',
]

Log = logging.getLogger('aiosqlite')

class Cursor:
    def __init__(
        self,
        conn: 'Connection',
        cursor: sqlite3.Cursor,
    ) -> None:
        self._conn = conn
        self._cursor = cursor

class Connection:
    def __init__(
        self,
        connector: Callable[[], sqlite3.Connection],
        loop: asyncio.AbstractEventLoop,
        executor: ThreadPoolExecutor,
    ) -> None:
        self._conn = None  # type: sqlite3.Connection
        self._connector = connector
        self._loop = loop
        self._executor = executor

    async def _execute(self, fn, *args, **kwargs):
        """Execute a function with the given arguments on the shared thread."""
        pt = partial(fn, *args, **kwargs)
        return await self._loop.run_in_executor(self._executor, pt)

    async def _connect(self):
        """Connect to the actual sqlite database."""
        if self._conn is None:
            self._conn = await self._execute(self._connector)

    async def __aenter__(self) -> 'Connection':
        await self._connect()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        await self.close()
        self._conn = None

    async def cursor(self) -> Cursor:
        raise NotImplementedError('Not yet available in aiosqlite')

    async def commit(self) -> None:
        raise NotImplementedError('Not yet available in aiosqlite')

    async def rollback(self) -> None:
        raise NotImplementedError('Not yet available in aiosqlite')

    async def close(self) -> None:
        await self._execute(self._conn.close)

    async def execute(
        self,
        sql: str,
        parameters: Iterable[Any] = None,
    ) -> Cursor:
        raise NotImplementedError('Not yet available in aiosqlite')

    async def executemany(
        self,
        sql: str,
        parameters: Iterable[Iterable[Any]] = None,
    ) -> Cursor:
        raise NotImplementedError('Not yet available in aiosqlite')

    async def executescript(
        self,
        sql_script: str,
    ) -> Cursor:
        raise NotImplementedError('Not yet available in aiosqlite')

def connect(
    database: str,
    **kwargs,
) -> Connection:
    """Create and return a connection proxy to the sqlite database."""

    loop = asyncio.get_event_loop()
    executor = ThreadPoolExecutor(1)

    def connector() -> sqlite3.Connection:
        return sqlite3.connect(
            database,
            check_same_thread=False,
            **kwargs,
        )

    return Connection(connector, loop, executor)
<commit_msg>Actual fix for py3.5<commit_after># Copyright 2017 John Reese
# Licensed under the MIT license

import asyncio
import logging
import sqlite3

from concurrent.futures import ThreadPoolExecutor
from functools import partial
from typing import Any, Callable, Iterable

__all__ = [
    'connect',
    'Connection',
    'Cursor',
]

Log = logging.getLogger('aiosqlite')

class Cursor:
    def __init__(
        self,
        conn: 'Connection',
        cursor: sqlite3.Cursor,
    ) -> None:
        self._conn = conn
        self._cursor = cursor

class Connection:
    def __init__(
        self,
        connector: Callable[[], sqlite3.Connection],
        loop: asyncio.AbstractEventLoop,
        executor: ThreadPoolExecutor,
    ) -> None:
        self._conn = None  # type: sqlite3.Connection
        self._connector = connector
        self._loop = loop
        self._executor = executor

    async def _execute(self, fn, *args, **kwargs):
        """Execute a function with the given arguments on the shared thread."""
        pt = partial(fn, *args, **kwargs)
        return await self._loop.run_in_executor(self._executor, pt)

    async def _connect(self):
        """Connect to the actual sqlite database."""
        if self._conn is None:
            self._conn = await self._execute(self._connector)

    async def __aenter__(self) -> 'Connection':
        await self._connect()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        await self.close()
        self._conn = None

    async def cursor(self) -> Cursor:
        raise NotImplementedError('Not yet available in aiosqlite')

    async def commit(self) -> None:
        raise NotImplementedError('Not yet available in aiosqlite')

    async def rollback(self) -> None:
        raise NotImplementedError('Not yet available in aiosqlite')

    async def close(self) -> None:
        await self._execute(self._conn.close)

    async def execute(
        self,
        sql: str,
        parameters: Iterable[Any] = None,
    ) -> Cursor:
        raise NotImplementedError('Not yet available in aiosqlite')

    async def executemany(
        self,
        sql: str,
        parameters: Iterable[Iterable[Any]] = None,
    ) -> Cursor:
        raise NotImplementedError('Not yet available in aiosqlite')

    async def executescript(
        self,
        sql_script: str,
    ) -> Cursor:
        raise NotImplementedError('Not yet available in aiosqlite')

def connect(
    database: str,
    **kwargs: Any
) -> Connection:
    """Create and return a connection proxy to the sqlite database."""

    loop = asyncio.get_event_loop()
    executor = ThreadPoolExecutor(1)

    def connector() -> sqlite3.Connection:
        return sqlite3.connect(
            database,
            check_same_thread=False,
            **kwargs,
        )

    return Connection(connector, loop, executor)
<|endoftext|>


<commit_before><commit_msg>Adding audioprocessor class<commit_after><|endoftext|>


<commit_before>"""
Toon van Eneco Thermostat Support.

This provides a component for the rebranded Quby thermostat as provided by
Eneco.
"""

import logging
from homeassistant.components.climate import (ClimateDevice,
                                              ATTR_TEMPERATURE,
                                              STATE_AUTO,
                                              STATE_HEAT,
                                              STATE_ECO,
                                              STATE_COOL,
                                              STATE_FAN_ONLY,
                                              SUPPORT_TARGET_TEMPERATURE,
                                              SUPPORT_OPERATION_MODE)
from homeassistant.const import (TEMP_CELSIUS, STATE_OFF)

import custom_components.toon as toon_main

_LOGGER = logging.getLogger(__name__)
SUPPORT_FLAGS = SUPPORT_TARGET_TEMPERATURE | SUPPORT_OPERATION_MODE

HA_TOON = {
                  STATE_AUTO: 'Comfort',
                  STATE_HEAT: 'Home',
                  STATE_ECO: 'Away',
                  STATE_COOL: 'Sleep'
                 }

TOON_HA = {value: key for key, value in HA_TOON.items()}

def setup_platform(hass, config, add_devices, discovery_info=None):
    """Setup thermostat."""
    # Add toon
    add_devices((ThermostatDevice(hass), ), True)

class ThermostatDevice(ClimateDevice):
    """Interface class for the toon module and HA."""

    def __init__(self, hass):
        """Initialize the device."""
        self._name = 'Toon van Eneco'
        self.hass = hass
        self.thermos = hass.data[toon_main.TOON_HANDLE]

        # set up internal state vars
        self._state = None
        self._temperature = None
        self._setpoint = None
        self._operation_list = [STATE_AUTO,
                                STATE_HEAT,
                                STATE_ECO,
                                STATE_COOL,
                                STATE_OFF,
                                STATE_FAN_ONLY]

    @property
    def supported_features(self):
        """Return the list of supported features."""
        return SUPPORT_FLAGS

    @property
    def name(self):
        """Name of this Thermostat."""
        return self._name

    @property
    def should_poll(self):
        """Polling is required."""
        return True

    @property
    def temperature_unit(self):
        """The unit of measurement used by the platform."""
        return TEMP_CELSIUS

    @property
    def current_operation(self):
        """Return current operation i.e. comfort, home, away."""
        if TOON_HA.get(self.thermos.get_data('state')) == None:
            return STATE_OFF
        else:
            return TOON_HA.get(self.thermos.get_data('state'))

    @property
    def operation_list(self):
        """List of available operation modes."""
        return self._operation_list

    @property
    def current_temperature(self):
        """Return the current temperature."""
        return self.thermos.get_data('temp')

    @property
    def target_temperature(self):
        """Return the temperature we try to reach."""
        return self.thermos.get_data('setpoint')

    def set_temperature(self, **kwargs):
        """Change the setpoint of the thermostat."""
        temp = kwargs.get(ATTR_TEMPERATURE)
        self.thermos.set_temp(temp)

    def set_operation_mode(self, operation_mode):
        """Set new operation mode as toonlib requires it."""
        self.thermos.set_state(HA_TOON[operation_mode])

    def update(self):
        """Update local state."""
        self.thermos.update()
<commit_msg>removed [fan only] mode<commit_after>"""
Toon van Eneco Thermostat Support.

This provides a component for the rebranded Quby thermostat as provided by
Eneco.
"""

import logging
from homeassistant.components.climate import (ClimateDevice,
                                              ATTR_TEMPERATURE,
                                              STATE_AUTO,
                                              STATE_HEAT,
                                              STATE_ECO,
                                              STATE_COOL,
                                              SUPPORT_TARGET_TEMPERATURE,
                                              SUPPORT_OPERATION_MODE)
from homeassistant.const import (TEMP_CELSIUS, STATE_OFF)

import custom_components.toon as toon_main

_LOGGER = logging.getLogger(__name__)
SUPPORT_FLAGS = SUPPORT_TARGET_TEMPERATURE | SUPPORT_OPERATION_MODE

HA_TOON = {
                  STATE_AUTO: 'Comfort',
                  STATE_HEAT: 'Home',
                  STATE_ECO: 'Away',
                  STATE_COOL: 'Sleep'
                 }

TOON_HA = {value: key for key, value in HA_TOON.items()}

def setup_platform(hass, config, add_devices, discovery_info=None):
    """Setup thermostat."""
    # Add toon
    add_devices((ThermostatDevice(hass), ), True)

class ThermostatDevice(ClimateDevice):
    """Interface class for the toon module and HA."""

    def __init__(self, hass):
        """Initialize the device."""
        self._name = 'Toon van Eneco'
        self.hass = hass
        self.thermos = hass.data[toon_main.TOON_HANDLE]

        # set up internal state vars
        self._state = None
        self._temperature = None
        self._setpoint = None
        self._operation_list = [STATE_AUTO,
                                STATE_HEAT,
                                STATE_ECO,
                                STATE_COOL,
                                STATE_OFF]

    @property
    def supported_features(self):
        """Return the list of supported features."""
        return SUPPORT_FLAGS

    @property
    def name(self):
        """Name of this Thermostat."""
        return self._name

    @property
    def should_poll(self):
        """Polling is required."""
        return True

    @property
    def temperature_unit(self):
        """The unit of measurement used by the platform."""
        return TEMP_CELSIUS

    @property
    def current_operation(self):
        """Return current operation i.e. comfort, home, away."""
        if TOON_HA.get(self.thermos.get_data('state')) == None:
            return STATE_OFF
        else:
            return TOON_HA.get(self.thermos.get_data('state'))

    @property
    def operation_list(self):
        """List of available operation modes."""
        return self._operation_list

    @property
    def current_temperature(self):
        """Return the current temperature."""
        return self.thermos.get_data('temp')

    @property
    def target_temperature(self):
        """Return the temperature we try to reach."""
        return self.thermos.get_data('setpoint')

    def set_temperature(self, **kwargs):
        """Change the setpoint of the thermostat."""
        temp = kwargs.get(ATTR_TEMPERATURE)
        self.thermos.set_temp(temp)

    def set_operation_mode(self, operation_mode):
        """Set new operation mode as toonlib requires it."""
        self.thermos.set_state(HA_TOON[operation_mode])

    def update(self):
        """Update local state."""
        self.thermos.update()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import numpy as np

def q(i,grid):
	"""This function returns a candidate state given the state that the Markov chain is currently in.

	Example
	-------
	n = (0,1,2)
	i = 1
	candidate = q(i)
	print(candidate)

	Parameters
	----------
	i : M x M array
		The current state of the Markov distribution (a gonnected graph).
	grid : M x 2 array
		x and y distance values for the vertices in the initial Markov state.

	Returns
	-------
	candidate : array
		candidate is the candidate state based on the decision calculus intrinsic to q."""

	i = np.array(i)
	grid = np.array(grid)
	candidate = i[:]
	rand_i = np.random.randint(len(i))
	rand_j = np.random.randint(len(i))

	while rand_i == rand_j: # or is_connected(rand_i,rand_j) == false:
	    rand_i = np.random.randint(len(i))
	    rand_j = np.random.randint(len(i))

	if candidate[rand_i,rand_j] != 0:
	    candidate[rand_i,rand_j] = 0
	    candidate[rand_j,rand_i] = 0

	elif candidate[rand_j,rand_i] == 0:
	    weight_i_j = np.sqrt((grid[rand_i,0] - grid[rand_j,0])**2 + (grid[rand_i,1]-grid[rand_j,1])**2)
	    candidate[rand_i,rand_j] = weight_i_j
	    candidate[rand_j,rand_i] = weight_i_j

	return candidate

def theta(i,grid,r=1):
	"""This function calculates the value of theta for the given state. Theta is a function defined in the problem statement.

	Example
	-------

	i = [[1,0],[0,1]]
	weights = [[1,2],[3,4]]
	value = theta(i,weights)

	Parameters
	----------

	i : MxM array
		This is the current state represented in an MxM adjacency matrix.

	grid : Mx2 array
		This is the array of x and y values in Cartesian space for the given index/label (i.e. the individual column headers in the 'i array.')

	Output
	------

	value : integer
		This is the theta value for the given state."""

	total_weight = np.sum(np.array(i))

def dijkstra(i):
	"""An implementation of Dijkstra's shortest path algorithm.
	Used pseudocode from http://www.gitta.info/Accessibiliti/en/html/Dijkstra_learningObject1.html."""
	import math
	dist = [0]
	previous = [0]
	for v in range(1,len(i),1):
	    dist.append(math.inf)
	    previous.append(0)

	Q = []
	for j in range(0,len(i),1):
	    Q.append(j)

	while len(Q) != 0:
	    dist_u = min(dist)
	    u = dist.index(min(dist))
	    Q.remove(u)
	    alt = 0
	    for v in Q:
	        alt = dist_u + i[u,v]
	        if alt < dist[v]:
	            dist[v] = alt
	            previous[v] = u
	return previous

def probability(i,j):
	"""This function computes the probability that the candidate state will be selected.

	Example
	-------

	i = [[1,0],[0,1]]
	j = [[0,1],[1,0]]
	prob = probability(i,j)

	Parameters
	----------

	i : MxM array
		This is the current graph represented by its adjacency matrix.

	j : MxM array
		This is the candidate graph represented by its adjacency matrix.

	Output
	------

	prob : float
		This is the probability that the candidate graph is selected. It takes values on [0,1]."""

def next_state(i,j,probability):
	""" This function returns the next state using functions already defined.

	Example
	-------

	state = next_state(2,3,.4)

	Parameters
	----------

	i : integer
		the current state

	j : integer
		the candidate state

	probability : float
		the probability that state j is selected

	Output
	------

	state : integer
		the state chosen"""

def iterator(i,eq_distrib,N):
	"""This is the main body of the implemented algorithm. This ignores the burn-off at the beginning of the simulation.

	Example
	-------
	n = (0,1,2)
	eq_distrib = (.4,.5,.1)
	N = 1000
	i = 2
	X, P = iterator(i,eq_distrib,N)
	print(X)
	print(P)

	Parameters
	----------

	i : integer
		the first state of the particle

	eq_distrib : 1-D array
		the equilibrium probability distribution of the n states in the state space

	N : the number of iterations desired

	Output
	------

	X : array of integers
		the record of all the states the Markov chain took
	P : nxn array of floats
		the transition probability matrix.
		"""

	import numpy as np

<commit_msg>wrote theta outline and probability function<commit_after># -*- coding: utf-8 -*-
import numpy as np

def q(i,grid):
	"""This function returns a candidate state given the state that the Markov chain is currently in.

	Example
	-------
	n = (0,1,2)
	i = 1
	candidate = q(i)
	print(candidate)

	Parameters
	----------
	i : M x M array
		The current state of the Markov distribution (a gonnected graph).
	grid : M x 2 array
		x and y distance values for the vertices in the initial Markov state.

	Returns
	-------
	candidate : array
		candidate is the candidate state based on the decision calculus intrinsic to q."""

	i = np.array(i)
	grid = np.array(grid)
	candidate = i[:]
	rand_i = np.random.randint(len(i))
	rand_j = np.random.randint(len(i))

	while rand_i == rand_j: # or is_connected(rand_i,rand_j) == false:
	    rand_i = np.random.randint(len(i))
	    rand_j = np.random.randint(len(i))

	if candidate[rand_i,rand_j] != 0:
	    candidate[rand_i,rand_j] = 0
	    candidate[rand_j,rand_i] = 0

	elif candidate[rand_j,rand_i] == 0:
	    weight_i_j = np.sqrt((grid[rand_i,0] - grid[rand_j,0])**2 + (grid[rand_i,1]-grid[rand_j,1])**2)
	    candidate[rand_i,rand_j] = weight_i_j
	    candidate[rand_j,rand_i] = weight_i_j

	return candidate

def theta(i,r=1):
	"""This function calculates the value of theta for the given state. Theta is a function defined in the problem statement.

	Example
	-------

	i = [[1,0],[0,1]]
	weights = [[1,2],[3,4]]
	value = theta(i,weights)

	Parameters
	----------

	i : MxM array
		This is the current state represented in an MxM adjacency matrix.

	grid : Mx2 array
		This is the array of x and y values in Cartesian space for the given index/label (i.e. the individual column headers in the 'i array.')

	Output
	------

	value : integer
		This is the theta value for the given state."""

	total_weight = np.sum(np.array(i))
	partial_weight = 0
	for v in len(i):
		path = dijkstra(np.array(i),v)
		for node in range(0,len(path)-1,1):
			partial_weight += i[path[node],path[node+1]]
	return(r * total_weight + partial_weight)

def dijkstra(i,desired_node):
	"""An implementation of Dijkstra's shortest path algorithm.
	Used pseudocode from http://www.gitta.info/Accessibiliti/en/html/Dijkstra_learningObject1.html."""
	import math
	dist = [0]
	previous = [0]
	for v in range(1,len(i),1):
	    dist.append(math.inf)
	    previous.append(0)

	Q = []
	for j in range(0,len(i),1):
	    Q.append(j)

	while len(Q) != 0:
	    dist_u = min(dist)
	    u = dist.index(min(dist))
	    Q.remove(u)
	    alt = 0
	    for v in Q:
	        alt = dist_u + i[u,v]
	        if alt < dist[v]:
	            dist[v] = alt
	            previous[v] = u
	return previous

def probability(i,j,T=1):
	"""This function computes the probability that the candidate state will be selected.

	Example
	-------

	i = [[1,0],[0,1]]
	j = [[0,1],[1,0]]
	prob = probability(i,j)

	Parameters
	----------

	i : MxM array
		This is the current graph represented by its adjacency matrix.

	j : MxM array
		This is the candidate graph represented by its adjacency matrix.

	T : float
		This is a parameter that can be tuned to improve the MCMC efficiency.

	Output
	------

	prob : float
		This is the probability that the candidate graph is selected. It takes values on [0,1]."""

	theta_i = theta(i)
	theta_j = theta(j)
	return(np.exp(-(theta_j - theta_i)/T))

def next_state(i,j,probability):
	""" This function returns the next state using functions already defined.

	Example
	-------

	state = next_state(2,3,.4)

	Parameters
	----------

	i : integer
		the current state

	j : integer
		the candidate state

	probability : float
		the probability that state j is selected

	Output
	------

	state : integer
		the state chosen"""

def iterator(i,eq_distrib,N):
	"""This is the main body of the implemented algorithm. This ignores the burn-off at the beginning of the simulation.

	Example
	-------
	n = (0,1,2)
	eq_distrib = (.4,.5,.1)
	N = 1000
	i = 2
	X, P = iterator(i,eq_distrib,N)
	print(X)
	print(P)

	Parameters
	----------

	i : integer
		the first state of the particle

	eq_distrib : 1-D array
		the equilibrium probability distribution of the n states in the state space

	N : the number of iterations desired

	Output
	------

	X : array of integers
		the record of all the states the Markov chain took
	P : nxn array of floats
		the transition probability matrix.
		"""

	import numpy as np

<|endoftext|>