<commit_before>#!/usr/bin/env python
# flake8: noqa: E402
""" Meant to be used by mesos-slave instead of the /usr/bin/docker executable
directly This will parse the CLI arguments intended for docker, extract
environment variable settings related to the actual node hostname and mesos
task ID, and use those as an additional --hostname argument when calling the
underlying docker command.

If the environment variables are unspecified, or if --hostname is already
specified, this does not change any arguments and just directly calls docker
as-is.

Additionally this wrapper will look for the environment variable
PIN_TO_NUMA_NODE which contains the physical CPU and memory to restrict the
container to. If the system is NUMA enabled, docker will be called with the
arguments cpuset-cpus and cpuset-mems.
"""
import logging
import os
import re
import socket
import sys

if "PATH" not in os.environ:
    # This command is sometimes executed in a sanitized environment
    # which does not have the path, which causes the following imports
    # to fail.
    # To compensate, we set a minimal path to get off the ground.
    os.environ["PATH"] = "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"

from paasta_tools.firewall import DEFAULT_SYNAPSE_SERVICE_DIR
from paasta_tools.firewall import firewall_flock
from paasta_tools.firewall import prepare_new_container
from paasta_tools.mac_address import reserve_unique_mac_address
from paasta_tools.utils import DEFAULT_SOA_DIR

LOCK_DIRECTORY = "/var/lib/paasta/mac-address"
ENV_MATCH_RE = re.compile(r"^(-\w*e\w*|--env(?P<file>-file)?)(=(?P<arg>\S.*))?$")
MAX_HOSTNAME_LENGTH = 60

def parse_env_args(args):
    result = dict(os.environ.items())
    in_env = False
    in_file = False
    for arg in args:
        if not in_env:
            match = ENV_MATCH_RE.match(arg)
            if not match:
                continue
            arg = match.group("arg") or ""
            in_file = bool(match.group("file"))
            if not arg:
                in_env = True
                continue

        in_env = False

        if in_file:
            result.update(read_env_file(arg))
            in_file = False
            continue

        try:
            k, v = arg.split("=", 1)
        except ValueError:
            continue

        result[k] = v

    return result

def read_env_file(filename):
    # Parse a file where each line is KEY=VALUE
    # return contents in dictionary form
    result = {}
    with open(filename) as f:
        for line in f:
            try:
                k, v = line.split("=", 1)
            except ValueError:
                continue
            result[k] = v.strip()
    return result

def can_add_hostname(args):
    # return False if --hostname is already specified or if --network=host
    if is_network_host(args):
        return False

    for index, arg in enumerate(args):

        # Check for --hostname and variants
        if arg == "-h":
            return False
        if arg.startswith("--hostname"):
            return False
        if len(arg) > 1 and arg[0] == "-" and arg[1] != "-":
            # several short args
            arg = arg.partition("=")[0]
            if "h" in arg:
                return False

    return True

def is_network_host(args):
    for index, arg in enumerate(args):
        # Check for --network=host and variants
        if arg in ("--net=host", "--network=host"):
            return True
        try:
            if arg in ("--net", "--network") and args[index + 1] == "host":
                return True
        except IndexError:
            pass

    return False

def is_run(args):
    try:
        list(args).index("run")
        return True
    except ValueError:
        return False

def can_add_mac_address(args):
    # return False if --mac-address is already specified or if --network=host
    if is_network_host(args) or not is_run(args):
        return False

    for index, arg in enumerate(args):
        # Check for --mac-address
        if arg.startswith("--mac-address"):
            return False

    return True

def generate_hostname_task_id(hostname, mesos_task_id):
    task_id = mesos_task_id.rpartition(".")[2]

    hostname_task_id = hostname + "-" + task_id

    # hostnames can only contain alphanumerics and dashes and must be no more
    # than 60 characters
    hostname_task_id = re.sub("[^a-zA-Z0-9-]+", "-", hostname_task_id)[
        :MAX_HOSTNAME_LENGTH
    ]

    # hostnames can also not end with dashes as per RFC952
    hostname_task_id = hostname_task_id.rstrip("-")

    return hostname_task_id

def add_argument(args, argument):
    # Add an argument immediately after 'run' command if it exists
    args = list(args)
    try:
        run_index = args.index("run")
    except ValueError:
        pass
    else:
        args.insert(run_index + 1, argument)
    return args

def get_cpumap():
    # Return a dict containing the core numbers per physical CPU
    core = 0
    cpumap = {}
    try:
        with open("/proc/cpuinfo", "r") as f:
            for line in f:
                m = re.match(r"physical\sid.*(\d)", line)
                if m:
                    cpuid = int(m.group(1))
                    if cpuid not in cpumap:
                        cpumap[cpuid] = []
                    cpumap[cpuid].append(core)
                    core += 1
    except IOError:
        logging.warning("Error while trying to read cpuinfo")
        pass
    return cpumap

def get_numa_memsize(nb_nodes):
    # Return memory size in mB per NUMA node assuming memory is split evenly
    # TODO: calculate and return real memory map
    try:
        with open("/proc/meminfo", "r") as f:
            for line in f:
                m = re.match(r"MemTotal:\s*(\d+)\skB", line)
                if m:
                    return int(m.group(1)) / 1024 / int(nb_nodes)
    except IOError:
        logging.warning("Error while trying to read meminfo")
        pass
    return 0

def arg_collision(new_args, current_args):
    # Returns True if one of the new arguments is already in the
    # current argument list.
    cur_arg_keys = []
    for c in current_args:
        cur_arg_keys.append(c.split("=")[0])
    return bool(set(new_args).intersection(set(cur_arg_keys)))

def is_numa_enabled():
    if os.path.exists("/proc/1/numa_maps"):
        return True
    else:
        logging.warning("The system does not support NUMA")
        return False

def get_cpu_requirements(env_args):
    # Ensure we return a float. If no requirements we return 0.0
    try:
        return float(env_args.get("MARATHON_APP_RESOURCE_CPUS"))
    except (ValueError, TypeError):
        logging.warning(
            "Could not read {} as a float".format(
                env_args.get("MARATHON_APP_RESOURCE_CPUS")
            )
        )
        return 0.0

def get_mem_requirements(env_args):
    # Ensure we return a float. If no requirements we return 0.0
    try:
        return float(env_args.get("MARATHON_APP_RESOURCE_MEM"))
    except (ValueError, TypeError):
        logging.warning(
            "Could not read {} as a float".format(
                env_args.get("MARATHON_APP_RESOURCE_MEM")
            )
        )
        return 0.0

def append_cpuset_args(argv, env_args):
    # Enable log messages to stderr
    logging.basicConfig(stream=sys.stderr, level=logging.INFO)

    try:
        pinned_numa_node = int(env_args.get("PIN_TO_NUMA_NODE"))
    except (ValueError, TypeError):
        logging.error(
            "Could not read PIN_TO_NUMA_NODE value as an int: {}".format(
                env_args.get("PIN_TO_NUMA_NODE")
            )
        )
        return argv

    cpumap = get_cpumap()

    if len(cpumap) < 1:
        logging.error("Less than 2 physical CPU detected")
        return argv
    if pinned_numa_node not in cpumap:
        logging.error(
            "Specified NUMA node: {} does not exist on this system".format(
                pinned_numa_node
            )
        )
        return argv
    if arg_collision(["--cpuset-cpus", "--cpuset-mems"], argv):
        logging.error("--cpuset options are already set. Not overriding")
        return argv
    if not is_numa_enabled():
        logging.error("Could not detect NUMA on the system")
        return argv
    if len(cpumap[pinned_numa_node]) < get_cpu_requirements(env_args):
        logging.error("The NUMA node has less cores than requested")
        return argv
    if get_numa_memsize(len(cpumap)) <= get_mem_requirements(env_args):
        logging.error(
            "Requested memory:{} MB does not fit in one NUMA node: {} MB".format(
                get_mem_requirements(env_args), get_numa_memsize(len(cpumap))
            )
        )
        return argv

    logging.info(f"Binding container to NUMA node {pinned_numa_node}")
    argv = add_argument(
        argv, ("--cpuset-cpus=" + ",".join(str(c) for c in cpumap[pinned_numa_node]))
    )
    argv = add_argument(argv, ("--cpuset-mems=" + str(pinned_numa_node)))
    return argv

def add_firewall(argv, service, instance):
    output = ""
    try:
        mac_address, lockfile = reserve_unique_mac_address(LOCK_DIRECTORY)
    except Exception as e:
        output = f"Unable to add mac address: {e}"
    else:
        argv = add_argument(argv, f"--mac-address={mac_address}")
        try:

            with firewall_flock():
                prepare_new_container(
                    DEFAULT_SOA_DIR,
                    DEFAULT_SYNAPSE_SERVICE_DIR,
                    service,
                    instance,
                    mac_address,
                )
        except Exception as e:
            output = f"Unable to add firewall rules: {e}"

    if output:
        print(output, file=sys.stderr)

    return argv

def main(argv=None):
    argv = argv if argv is not None else sys.argv

    env_args = parse_env_args(argv)

    if env_args.get("PIN_TO_NUMA_NODE"):
        argv = append_cpuset_args(argv, env_args)

    # Marathon sets MESOS_TASK_ID
    mesos_task_id = env_args.get("MESOS_TASK_ID")

    if mesos_task_id and can_add_hostname(argv):
        hostname = socket.getfqdn()
        argv = add_argument(argv, f"-e=PAASTA_HOST={hostname}")
        hostname_task_id = generate_hostname_task_id(
            hostname.partition(".")[0], mesos_task_id
        )
        argv = add_argument(argv, f"--hostname={hostname_task_id }")

    paasta_firewall = env_args.get("PAASTA_FIREWALL")
    service = env_args.get("PAASTA_SERVICE")
    instance = env_args.get("PAASTA_INSTANCE")
    if paasta_firewall and service and instance and can_add_mac_address(argv):
        try:
            argv = add_firewall(argv, service, instance)
        except Exception as e:
            print(f"Unhandled exception in add_firewall: {e}", file=sys.stderr)

    os.execlp("docker", "docker", *argv[1:])
<commit_msg>Delay imports until necessary<commit_after>#!/usr/bin/env python
# flake8: noqa: E402
""" Meant to be used by mesos-slave instead of the /usr/bin/docker executable
directly This will parse the CLI arguments intended for docker, extract
environment variable settings related to the actual node hostname and mesos
task ID, and use those as an additional --hostname argument when calling the
underlying docker command.

If the environment variables are unspecified, or if --hostname is already
specified, this does not change any arguments and just directly calls docker
as-is.

Additionally this wrapper will look for the environment variable
PIN_TO_NUMA_NODE which contains the physical CPU and memory to restrict the
container to. If the system is NUMA enabled, docker will be called with the
arguments cpuset-cpus and cpuset-mems.
"""
import logging
import os
import re
import socket
import sys

if "PATH" not in os.environ:
    # This command is sometimes executed in a sanitized environment
    # which does not have the path, which causes the following imports
    # to fail.
    # To compensate, we set a minimal path to get off the ground.
    os.environ["PATH"] = "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"

LOCK_DIRECTORY = "/var/lib/paasta/mac-address"
ENV_MATCH_RE = re.compile(r"^(-\w*e\w*|--env(?P<file>-file)?)(=(?P<arg>\S.*))?$")
MAX_HOSTNAME_LENGTH = 60

def parse_env_args(args):
    result = dict(os.environ.items())
    in_env = False
    in_file = False
    for arg in args:
        if not in_env:
            match = ENV_MATCH_RE.match(arg)
            if not match:
                continue
            arg = match.group("arg") or ""
            in_file = bool(match.group("file"))
            if not arg:
                in_env = True
                continue

        in_env = False

        if in_file:
            result.update(read_env_file(arg))
            in_file = False
            continue

        try:
            k, v = arg.split("=", 1)
        except ValueError:
            continue

        result[k] = v

    return result

def read_env_file(filename):
    # Parse a file where each line is KEY=VALUE
    # return contents in dictionary form
    result = {}
    with open(filename) as f:
        for line in f:
            try:
                k, v = line.split("=", 1)
            except ValueError:
                continue
            result[k] = v.strip()
    return result

def can_add_hostname(args):
    # return False if --hostname is already specified or if --network=host
    if is_network_host(args):
        return False

    for index, arg in enumerate(args):

        # Check for --hostname and variants
        if arg == "-h":
            return False
        if arg.startswith("--hostname"):
            return False
        if len(arg) > 1 and arg[0] == "-" and arg[1] != "-":
            # several short args
            arg = arg.partition("=")[0]
            if "h" in arg:
                return False

    return True

def is_network_host(args):
    for index, arg in enumerate(args):
        # Check for --network=host and variants
        if arg in ("--net=host", "--network=host"):
            return True
        try:
            if arg in ("--net", "--network") and args[index + 1] == "host":
                return True
        except IndexError:
            pass

    return False

def is_run(args):
    try:
        list(args).index("run")
        return True
    except ValueError:
        return False

def can_add_mac_address(args):
    # return False if --mac-address is already specified or if --network=host
    if is_network_host(args) or not is_run(args):
        return False

    for index, arg in enumerate(args):
        # Check for --mac-address
        if arg.startswith("--mac-address"):
            return False

    return True

def generate_hostname_task_id(hostname, mesos_task_id):
    task_id = mesos_task_id.rpartition(".")[2]

    hostname_task_id = hostname + "-" + task_id

    # hostnames can only contain alphanumerics and dashes and must be no more
    # than 60 characters
    hostname_task_id = re.sub("[^a-zA-Z0-9-]+", "-", hostname_task_id)[
        :MAX_HOSTNAME_LENGTH
    ]

    # hostnames can also not end with dashes as per RFC952
    hostname_task_id = hostname_task_id.rstrip("-")

    return hostname_task_id

def add_argument(args, argument):
    # Add an argument immediately after 'run' command if it exists
    args = list(args)
    try:
        run_index = args.index("run")
    except ValueError:
        pass
    else:
        args.insert(run_index + 1, argument)
    return args

def get_cpumap():
    # Return a dict containing the core numbers per physical CPU
    core = 0
    cpumap = {}
    try:
        with open("/proc/cpuinfo", "r") as f:
            for line in f:
                m = re.match(r"physical\sid.*(\d)", line)
                if m:
                    cpuid = int(m.group(1))
                    if cpuid not in cpumap:
                        cpumap[cpuid] = []
                    cpumap[cpuid].append(core)
                    core += 1
    except IOError:
        logging.warning("Error while trying to read cpuinfo")
        pass
    return cpumap

def get_numa_memsize(nb_nodes):
    # Return memory size in mB per NUMA node assuming memory is split evenly
    # TODO: calculate and return real memory map
    try:
        with open("/proc/meminfo", "r") as f:
            for line in f:
                m = re.match(r"MemTotal:\s*(\d+)\skB", line)
                if m:
                    return int(m.group(1)) / 1024 / int(nb_nodes)
    except IOError:
        logging.warning("Error while trying to read meminfo")
        pass
    return 0

def arg_collision(new_args, current_args):
    # Returns True if one of the new arguments is already in the
    # current argument list.
    cur_arg_keys = []
    for c in current_args:
        cur_arg_keys.append(c.split("=")[0])
    return bool(set(new_args).intersection(set(cur_arg_keys)))

def is_numa_enabled():
    if os.path.exists("/proc/1/numa_maps"):
        return True
    else:
        logging.warning("The system does not support NUMA")
        return False

def get_cpu_requirements(env_args):
    # Ensure we return a float. If no requirements we return 0.0
    try:
        return float(env_args.get("MARATHON_APP_RESOURCE_CPUS"))
    except (ValueError, TypeError):
        logging.warning(
            "Could not read {} as a float".format(
                env_args.get("MARATHON_APP_RESOURCE_CPUS")
            )
        )
        return 0.0

def get_mem_requirements(env_args):
    # Ensure we return a float. If no requirements we return 0.0
    try:
        return float(env_args.get("MARATHON_APP_RESOURCE_MEM"))
    except (ValueError, TypeError):
        logging.warning(
            "Could not read {} as a float".format(
                env_args.get("MARATHON_APP_RESOURCE_MEM")
            )
        )
        return 0.0

def append_cpuset_args(argv, env_args):
    # Enable log messages to stderr
    logging.basicConfig(stream=sys.stderr, level=logging.INFO)

    try:
        pinned_numa_node = int(env_args.get("PIN_TO_NUMA_NODE"))
    except (ValueError, TypeError):
        logging.error(
            "Could not read PIN_TO_NUMA_NODE value as an int: {}".format(
                env_args.get("PIN_TO_NUMA_NODE")
            )
        )
        return argv

    cpumap = get_cpumap()

    if len(cpumap) < 1:
        logging.error("Less than 2 physical CPU detected")
        return argv
    if pinned_numa_node not in cpumap:
        logging.error(
            "Specified NUMA node: {} does not exist on this system".format(
                pinned_numa_node
            )
        )
        return argv
    if arg_collision(["--cpuset-cpus", "--cpuset-mems"], argv):
        logging.error("--cpuset options are already set. Not overriding")
        return argv
    if not is_numa_enabled():
        logging.error("Could not detect NUMA on the system")
        return argv
    if len(cpumap[pinned_numa_node]) < get_cpu_requirements(env_args):
        logging.error("The NUMA node has less cores than requested")
        return argv
    if get_numa_memsize(len(cpumap)) <= get_mem_requirements(env_args):
        logging.error(
            "Requested memory:{} MB does not fit in one NUMA node: {} MB".format(
                get_mem_requirements(env_args), get_numa_memsize(len(cpumap))
            )
        )
        return argv

    logging.info(f"Binding container to NUMA node {pinned_numa_node}")
    argv = add_argument(
        argv, ("--cpuset-cpus=" + ",".join(str(c) for c in cpumap[pinned_numa_node]))
    )
    argv = add_argument(argv, ("--cpuset-mems=" + str(pinned_numa_node)))
    return argv

def add_firewall(argv, service, instance):
    # Delay importing these until we really need to, because they are
    # expensive to import and we need to be fast, most of the time
    from paasta_tools.firewall import DEFAULT_SYNAPSE_SERVICE_DIR
    from paasta_tools.firewall import firewall_flock
    from paasta_tools.firewall import prepare_new_container
    from paasta_tools.mac_address import reserve_unique_mac_address
    from paasta_tools.utils import DEFAULT_SOA_DIR

    output = ""
    try:
        mac_address, lockfile = reserve_unique_mac_address(LOCK_DIRECTORY)
    except Exception as e:
        output = f"Unable to add mac address: {e}"
    else:
        argv = add_argument(argv, f"--mac-address={mac_address}")
        try:

            with firewall_flock():
                prepare_new_container(
                    DEFAULT_SOA_DIR,
                    DEFAULT_SYNAPSE_SERVICE_DIR,
                    service,
                    instance,
                    mac_address,
                )
        except Exception as e:
            output = f"Unable to add firewall rules: {e}"

    if output:
        print(output, file=sys.stderr)

    return argv

def main(argv=None):
    argv = argv if argv is not None else sys.argv

    env_args = parse_env_args(argv)

    if env_args.get("PIN_TO_NUMA_NODE"):
        argv = append_cpuset_args(argv, env_args)

    # Marathon sets MESOS_TASK_ID
    mesos_task_id = env_args.get("MESOS_TASK_ID")

    if mesos_task_id and can_add_hostname(argv):
        hostname = socket.getfqdn()
        argv = add_argument(argv, f"-e=PAASTA_HOST={hostname}")
        hostname_task_id = generate_hostname_task_id(
            hostname.partition(".")[0], mesos_task_id
        )
        argv = add_argument(argv, f"--hostname={hostname_task_id }")

    paasta_firewall = env_args.get("PAASTA_FIREWALL")
    service = env_args.get("PAASTA_SERVICE")
    instance = env_args.get("PAASTA_INSTANCE")
    if paasta_firewall and service and instance and can_add_mac_address(argv):
        try:
            argv = add_firewall(argv, service, instance)
        except Exception as e:
            print(f"Unhandled exception in add_firewall: {e}", file=sys.stderr)

    os.execlp("docker", "docker", *argv[1:])
<|endoftext|>


<commit_before>"""Helper to output stuff"""

from tabulate import tabulate
import os

def print_list(header, entries):
    """Prints out a list"""
    print(tabulate(fix_descriptions(entries), header, tablefmt="presto"))

def print_table(entries):
    """Prints out a table"""
    print(tabulate(fix_descriptions(entries), tablefmt="plain"))

def std_input(text, default):
    """Get input or return default if none is given."""
    return input(text.format(default)) or default

def fix_descriptions(entries):
    """Fixes the description to fit into the terminal"""

    clean_entries = []
    ml = get_max_desc_width(get_longest_list(entries))

    for entry in entries:
        clean_entry = entry
        max_value = max(entry, key=len)
        for idx, val in enumerate(entry):
            if val is max_value:
                clean_entry[idx] = entry[idx][:ml] + (entry[idx][ml:] and ' [..]')
        clean_entries.append(clean_entry)

    return clean_entries

def get_longest_list(entries):
    max_list = ['']*len(entries[0])
    for entry in entries:
        for idx, val in enumerate(entry):
            if len(val) > len(max_list[idx]):
                max_list[idx] = val
    return max_list

def get_max_desc_width(lst):
    _, columns = os.popen('stty size', 'r').read().split()
    length = int(columns)
    max_value = max(lst, key=len)
    for val in lst:
        if val is not max_value:
            length -= len(val)

    return length - 15
<commit_msg>Fix cutting descriptions<commit_after>"""Helper to output stuff"""

from tabulate import tabulate
import os

def print_list(header, entries):
    """Prints out a list"""
    print(tabulate(fix_descriptions(entries), header, tablefmt="presto"))

def print_table(entries):
    """Prints out a table"""
    print(tabulate(fix_descriptions(entries), tablefmt="plain"))

def std_input(text, default):
    """Get input or return default if none is given."""
    return input(text.format(default)) or default

def fix_descriptions(entries):
    """Fixes the description to fit into the terminal"""

    clean_entries = []
    ml = get_max_desc_width(get_longest_list(entries))

    for entry in entries:
        clean_entry = entry
        max_value = max(entry, key=len)
        for idx, val in enumerate(entry):
            if val is max_value:
                clean_entry[idx] = entry[idx][:ml] + (entry[idx][ml:] and ' [..]')
        clean_entries.append(clean_entry)

    return clean_entries

def get_longest_list(entries):
    max_list = ['']*len(entries[0])
    for entry in entries:
        for idx, val in enumerate(entry):
            if len(val) > len(max_list[idx]):
                max_list[idx] = val
    return max_list

def get_max_desc_width(lst):
    _, columns = os.popen('stty size', 'r').read().split()
    length = int(columns)
    max_value = max(lst, key=len)
    for val in lst:
        if val is not max_value:
            length -= len(val)

    return length - 20
<|endoftext|>


<commit_before>"""
Module for comutation of densities of states for atom molecule collsions
"""
import numpy as np
import scipy.constants
from math import pi
from collections import OrderedDict
import yaml

# pylint: disable=E1103
# pylint: disable=R0902
# pylint: disable=R0903
# pylint: disable=R0913,R0914

# define physical constants
EH2K = scipy.constants.physical_constants["hartree-kelvin relationship"][0]
K2EH = scipy.constants.physical_constants["kelvin-hartree relationship"][0]
ATOMICUNITOFTIME = scipy.constants.physical_constants["atomic unit of time"][0]

def get_data(filen):
    """reads a file of the form 1st line heading columns
       2nd lines on values for those headings
       returns a nested dictionary
       converts values to float if possible"""

    with open(filen) as input:
        try:
            data_dict = yaml.load(input)
        except yaml.YAMLError as exc:
            print(exc)
#     open_file = open(filen, "r")
#     lines = open_file.readlines()
#     keys = lines[0].split()

#     data_dict = OrderedDict()

#     for line in lines[1:]:
#         items = line.split()
#         for i, item in enumerate(items):
#             try:
#                 items[i] = float(item)
#             except ValueError:
#                 pass
#         sub_dict = dict(zip(keys, items))
#         data_dict.update({sub_dict['sys']: sub_dict})
#     outname = filen.replace('txt','yml')
#     with open(outname, 'w') as outfile:
#             yaml.dump(dict(data_dict), outfile, default_flow_style=True)
    return data_dict

def _read_data(dir_loc, inputfile, qnmax):
    """read in the data from precalculated rovibrational energy levels of the
    system files are og the form 'txtjXXXtxt' where XXX is the L quantum number
    starting at 0
    """
    bound_states = []  # bound states for different L quntum numbers
    for i in xrange(qnmax+1):
        filen = dir_loc + inputfile.replace('j', 'j'+str(i))
        try:
            data = np.sort(np.loadtxt(filen, usecols=(1,)))
        except IOError:  # maximum quantum number reached set qnmax and return
            qnmax = i-1
            return bound_states, qnmax
        if len(np.atleast_1d(data)) == 1:
            data = np.reshape(data, 1)  # accounts for behavior of 0d arrays
        bound_states.append(data)
    return bound_states, qnmax

class AMDOS(object):
    """
    Atom molecule density of states class
    """
    def __init__(self, filen):
        self.systems = get_data(filen)
        self.bound_states_c = None
        self.bound_states_d = None
        self.lmax = None
        self.nmax = None
        self.lifetime = None
        self.dos = None
        self.num_open = None

    @staticmethod
    def _dos(jqn, mqn, bound_states_c, bound_states_d, vmax,
             energy_range=10.0, energy_offset=0.0):
        """
        computes the dos for an atom moleucle collsiions as described in
        Statistical aspects of ultracold resonant scattering Mayle, Ruzic and
        Bohn
        Variables
        jqn       - total jqn quantum number.
        mqn     - mj_qn projection quantum number.
        bound_states_c - boundstates of the complex obtained from _read_data
        bound_states_d - boundstates of the molecule obtained from _read_data
        vmax    - maximum vibrational quantum number.
        energy_range  - dos = number of states counted/energy range
        energy_offset - calculate dos around energy = energy_offset
        """
        if abs(mqn) > abs(jqn):
            raise Exception('physically impossible (abs(mj_qn) > abs(jqn))')
        limit = (energy_range/2.0)*K2EH  # convert energy_range K to hartrees
        energy_offset *= K2EH
        outside = np.ma.masked_outside  # mask enties in list outside range
        count = np.ma.count  # counts the number of unmasked entries in a list
        abs_gs = bound_states_d[0][0]
        num = 0
        # looping over rotational state of dimer a
        for nqn in xrange(len(bound_states_d)):
            # looping over all l constant with jqn
            for lqn in xrange(len(bound_states_c)):
                # only include pairs which couple to form jqn
                if abs(nqn-lqn) <= jqn and nqn+lqn >= jqn:
                    # looping over all vibrational levels of a
                    for vqn in xrange(min(len(bound_states_d[nqn]), vmax+1)):
                        # degeneracy
                        deg = len(xrange(max(-lqn, mqn-nqn),
                                         min(lqn, mqn+nqn)))+1
                        threshold_energy = bound_states_d[nqn][vqn]-abs_gs
                        num += count(
                            outside(bound_states_c[lqn],
                                    -limit-threshold_energy+energy_offset,
                                    limit-threshold_energy+energy_offset))*deg

        dos = (float(num)/energy_range)*1.0E-3
        lifetime = dos*1.0E3*EH2K*2.0*pi*ATOMICUNITOFTIME*1.0e9
        # return dos in mK-1 and lifetime in ns
        return dos, lifetime

    @staticmethod
    def _num_open(jqn, mqn, bound_states_c, bound_states_d,
                  vmax, energy_offset=0.0):
        """
        computes the number of threshold channels of the complex with energy
        less than energy_offset.
        Statistical aspects of ultracold resonant scattering Mayle, Rizic and
        Bohn
        VAriables
        jqn       - total J quantum number.
        mqn     - mj_qn projection quantum number.
        bound_states_c - boundstates of the complex obtained from _read_data
        bound_states_d - boundstates of the molecule obtained from _read_data
        vmax    - maximum vibrational quantum number.
        energy_offset - calculate number of threshold channels of the complex
                        with energy less than energy_offset.
        """
        if abs(mqn) > abs(jqn):
            print 'physically impossible (abs(mj_qn) > abs(jqn))'
            return
        energy_offset *= K2EH
        abs_gs = bound_states_d[0][0]
        num_open = 0
        # looping over rotational state of dimer a
        for nqn in xrange(len(bound_states_d)):
            # looping over all l constant with jqn
            for lqn in xrange(len(bound_states_c)):
                # include pairs of rotational levels which couple to form jqn
                if abs(nqn-lqn) <= jqn and nqn+lqn >= jqn:
                    # looping over all vibrational levels of a
                    for vqn in xrange(min(len(bound_states_d[nqn]), vmax+1)):
                        threshold_energy = bound_states_d[nqn][vqn]-abs_gs
                        if threshold_energy < energy_offset:
                            num_open += 1
        return num_open

    def get_dos(self, key, jqn=0, mqn=0, nmax=100, vmax=9999,
                energy_offset=0.0):
        """
        simple wrapper around _read_data and AMDOS
        Variables
        jqn       - total J quantum number.
        mqn     - mj_qn projection quantum number.
        nmax    - maximum allowed rotaional quantum number of the molecule.
        vmax    - maximum allowed vibrational quantum number of the molecule.
        energy_offset - calculate dos around energy = energy_offset
        """
        lmax = nmax + jqn
        self.bound_states_d, self.nmax = _read_data(
            self.systems[key]['dimer_dirn'],
            self.systems[key]['dimer_filen'],
            nmax)
        self.bound_states_c, self.lmax = _read_data(
            self.systems[key]['cmplx_dirn'],
            self.systems[key]['cmplx_filen'],
            lmax)
        self.dos, self.lifetime = AMDOS._dos(
            jqn,
            mqn,
            self.bound_states_c,
            self.bound_states_d,
            vmax,
            energy_offset=energy_offset)
        return self.dos, self.lifetime

    def get_num_open(self, key, jqn=0, mqn=0, nmax=100, vmax=9999,
                     energy_offset=0.0):
        """
        simple wrapper around _read_data and _num_open
        Variables
        jqn       - total jqn quantum number.
        mqn     - mj_qn projection quantum number.
        nmax    - maximum allowed rotaional quantum number of the molecule.
        vmax    - maximum allowed vibrational quantum number of the molecule.
        energy_offset - calculate dos around energy = energy_offset
        """
        lmax = nmax + jqn
        self.bound_states_d, self.nmax = _read_data(
            self.systems[key]['dimer_dirn'],
            self.systems[key]['dimer_filen'],
            nmax)
        self.bound_states_c, self.lmax = _read_data(
            self.systems[key]['cmplx_dirn'],
            self.systems[key]['cmplx_filen'],
            lmax)
        self.num_open = AMDOS._num_open(
            jqn,
            mqn,
            self.bound_states_c,
            self.bound_states_d,
            vmax,
            energy_offset)
        return self.num_open

class MMDOS(object):
    """
    class for molecule molecule density of states calculations
    """
    def __init__(self, filen):
        self.systems = get_data(filen)
        self.nmax = None
        self.lmax = None
        self.bound_states_c = None
        self.bound_states_d = None
        self.lifetime = None
        self.dos = None

    @staticmethod
    def _mm_dos(jqn, mj_qn, nmax, vmax, lmax,
                bound_states_d, bound_states_c,
                energy_range=10.0):
        """
        computes the dos for an atom moleucle collsiions as described in
        Scattering of Ultracold Molecules in the Highly Resonant Regime --
        Mayle, Ruzic, Quememer and Bohn
        Variables
        jqn       - total J quantum number.
        mqn     - mj_qn projection quantum number.
        bound_states_c - boundstates of the complex obtained from _read_data
        bound_states_d - boundstates of the molecule obtained from _read_data
        vmax    - maximum vibrational quantum number.
        lmax    - maximum end-over-end rotational quantum number of the two
                  molecules
        nmax    - maximum rotational quantum number of a single molecule.
        energy_range  - dos = number of states counted/energy range
        """
        if abs(mj_qn) > abs(jqn):
            raise Exception('physically impossible (abs(mj_qn) > abs(jqn))')
        num = 0  # variable to hold the number of states between limits
        limit = (energy_range/2.0)*K2EH  # convert energy_range K to hartrees
        abs_gs = bound_states_d[0][0]  # energy of the absolute ground state
        for nqn in xrange(max(0, jqn-lmax), min(2*nmax, lmax-jqn)+1):
            for lqn in xrange(abs(jqn-nqn), min(lmax, jqn+nqn)+1):
                for n1qn in xrange(max(0, nqn-nmax), nmax+1):
                    for n2qn in xrange(abs(nqn-n1qn), min(nqn+n1qn, nmax)+1):
                        # looping over all vibrational levels of dimer1
                        for v1qn in xrange(min(len(bound_states_d[n1qn]),
                                               vmax+1)):
                            # looping over all vibrational levels of dimer2
                            for v2qn in xrange(min(len(bound_states_d[n2qn]),
                                                   vmax+1)):
                                threshold_energy = (
                                    bound_states_d[n1qn][v1qn] - abs_gs +
                                    bound_states_d[n2qn][v2qn] - abs_gs)
                                if (bound_states_c[lqn][0] >
                                        limit-threshold_energy):
                                    # the lowest state is higher than highest
                                    # threshold
                                    break
                                else:
                                    start_end = np.searchsorted(
                                        bound_states_c[lqn],
                                        [-limit-threshold_energy,
                                         limit-threshold_energy],
                                        'left')
                                    deg = (abs(max(-lqn, mj_qn-nqn)) +
                                           min(lqn, mj_qn+nqn) + 1)
                                    num += deg*(start_end[1] - start_end[0])
        # return dos in per uK i.e.times 10**-6K
        dos = (float(num)/energy_range)*1.0E-6
        # return lifetime in ms
        lifetime = dos*1.0E6*EH2K*2.0*pi*ATOMICUNITOFTIME * 1.0e3
        return dos, lifetime

    def mm_dos(self, key, jqn=0, mqn=0, nmax=100, vmax=999):
        """
        simple wrapper around _read_data and MMDOS
        Variables
        jqn       - total J quantum number.
        mqn     - mj_qn projection quantum number.
        nmax    - maximum allowed rotaional quantum number of the molecule.
        vmax    - maximum allowed vibrational quantum number of the molecule.
        energy_offset - calculate dos around energy = energy_offset
        """
        lmax = 2*nmax+jqn
        self.bound_states_d, self.nmax = _read_data(
            self.systems[key]['dimer_dirn'],
            self.systems[key]['dimer_filen'],
            nmax)
        self.bound_states_c, self.lmax = _read_data(
            self.systems[key]['cmplx_dirn'],
            self.systems[key]['cmplx_filen'],
            lmax)
        self.dos, self.lifetime = MMDOS._mm_dos(
            jqn, mqn, self.nmax, vmax, self.lmax,
            self.bound_states_d, self.bound_states_c)
        return self.dos, self.lifetime
<commit_msg>remove commetns<commit_after>"""
Module for comutation of densities of states for atom molecule collsions
"""
import numpy as np
import scipy.constants
from math import pi
from collections import OrderedDict
import yaml

# pylint: disable=E1103
# pylint: disable=R0902
# pylint: disable=R0903
# pylint: disable=R0913,R0914

# define physical constants
EH2K = scipy.constants.physical_constants["hartree-kelvin relationship"][0]
K2EH = scipy.constants.physical_constants["kelvin-hartree relationship"][0]
ATOMICUNITOFTIME = scipy.constants.physical_constants["atomic unit of time"][0]

def get_data(filen):
    """reads a file of the form 1st line heading columns
       2nd lines on values for those headings
       returns a nested dictionary
       converts values to float if possible"""

    with open(filen) as input:
        try:
            data_dict = yaml.load(input)
        except yaml.YAMLError as exc:
            print(exc)
    return data_dict

def _read_data(dir_loc, inputfile, qnmax):
    """read in the data from precalculated rovibrational energy levels of the
    system files are og the form 'txtjXXXtxt' where XXX is the L quantum number
    starting at 0
    """
    bound_states = []  # bound states for different L quntum numbers
    for i in xrange(qnmax+1):
        filen = dir_loc + inputfile.replace('j', 'j'+str(i))
        try:
            data = np.sort(np.loadtxt(filen, usecols=(1,)))
        except IOError:  # maximum quantum number reached set qnmax and return
            qnmax = i-1
            return bound_states, qnmax
        if len(np.atleast_1d(data)) == 1:
            data = np.reshape(data, 1)  # accounts for behavior of 0d arrays
        bound_states.append(data)
    return bound_states, qnmax

class AMDOS(object):
    """
    Atom molecule density of states class
    """
    def __init__(self, filen):
        self.systems = get_data(filen)
        self.bound_states_c = None
        self.bound_states_d = None
        self.lmax = None
        self.nmax = None
        self.lifetime = None
        self.dos = None
        self.num_open = None

    @staticmethod
    def _dos(jqn, mqn, bound_states_c, bound_states_d, vmax,
             energy_range=10.0, energy_offset=0.0):
        """
        computes the dos for an atom moleucle collsiions as described in
        Statistical aspects of ultracold resonant scattering Mayle, Ruzic and
        Bohn
        Variables
        jqn       - total jqn quantum number.
        mqn     - mj_qn projection quantum number.
        bound_states_c - boundstates of the complex obtained from _read_data
        bound_states_d - boundstates of the molecule obtained from _read_data
        vmax    - maximum vibrational quantum number.
        energy_range  - dos = number of states counted/energy range
        energy_offset - calculate dos around energy = energy_offset
        """
        if abs(mqn) > abs(jqn):
            raise Exception('physically impossible (abs(mj_qn) > abs(jqn))')
        limit = (energy_range/2.0)*K2EH  # convert energy_range K to hartrees
        energy_offset *= K2EH
        outside = np.ma.masked_outside  # mask enties in list outside range
        count = np.ma.count  # counts the number of unmasked entries in a list
        abs_gs = bound_states_d[0][0]
        num = 0
        # looping over rotational state of dimer a
        for nqn in xrange(len(bound_states_d)):
            # looping over all l constant with jqn
            for lqn in xrange(len(bound_states_c)):
                # only include pairs which couple to form jqn
                if abs(nqn-lqn) <= jqn and nqn+lqn >= jqn:
                    # looping over all vibrational levels of a
                    for vqn in xrange(min(len(bound_states_d[nqn]), vmax+1)):
                        # degeneracy
                        deg = len(xrange(max(-lqn, mqn-nqn),
                                         min(lqn, mqn+nqn)))+1
                        threshold_energy = bound_states_d[nqn][vqn]-abs_gs
                        num += count(
                            outside(bound_states_c[lqn],
                                    -limit-threshold_energy+energy_offset,
                                    limit-threshold_energy+energy_offset))*deg

        dos = (float(num)/energy_range)*1.0E-3
        lifetime = dos*1.0E3*EH2K*2.0*pi*ATOMICUNITOFTIME*1.0e9
        # return dos in mK-1 and lifetime in ns
        return dos, lifetime

    @staticmethod
    def _num_open(jqn, mqn, bound_states_c, bound_states_d,
                  vmax, energy_offset=0.0):
        """
        computes the number of threshold channels of the complex with energy
        less than energy_offset.
        Statistical aspects of ultracold resonant scattering Mayle, Rizic and
        Bohn
        VAriables
        jqn       - total J quantum number.
        mqn     - mj_qn projection quantum number.
        bound_states_c - boundstates of the complex obtained from _read_data
        bound_states_d - boundstates of the molecule obtained from _read_data
        vmax    - maximum vibrational quantum number.
        energy_offset - calculate number of threshold channels of the complex
                        with energy less than energy_offset.
        """
        if abs(mqn) > abs(jqn):
            print 'physically impossible (abs(mj_qn) > abs(jqn))'
            return
        energy_offset *= K2EH
        abs_gs = bound_states_d[0][0]
        num_open = 0
        # looping over rotational state of dimer a
        for nqn in xrange(len(bound_states_d)):
            # looping over all l constant with jqn
            for lqn in xrange(len(bound_states_c)):
                # include pairs of rotational levels which couple to form jqn
                if abs(nqn-lqn) <= jqn and nqn+lqn >= jqn:
                    # looping over all vibrational levels of a
                    for vqn in xrange(min(len(bound_states_d[nqn]), vmax+1)):
                        threshold_energy = bound_states_d[nqn][vqn]-abs_gs
                        if threshold_energy < energy_offset:
                            num_open += 1
        return num_open

    def get_dos(self, key, jqn=0, mqn=0, nmax=100, vmax=9999,
                energy_offset=0.0):
        """
        simple wrapper around _read_data and AMDOS
        Variables
        jqn       - total J quantum number.
        mqn     - mj_qn projection quantum number.
        nmax    - maximum allowed rotaional quantum number of the molecule.
        vmax    - maximum allowed vibrational quantum number of the molecule.
        energy_offset - calculate dos around energy = energy_offset
        """
        lmax = nmax + jqn
        self.bound_states_d, self.nmax = _read_data(
            self.systems[key]['dimer_dirn'],
            self.systems[key]['dimer_filen'],
            nmax)
        self.bound_states_c, self.lmax = _read_data(
            self.systems[key]['cmplx_dirn'],
            self.systems[key]['cmplx_filen'],
            lmax)
        self.dos, self.lifetime = AMDOS._dos(
            jqn,
            mqn,
            self.bound_states_c,
            self.bound_states_d,
            vmax,
            energy_offset=energy_offset)
        return self.dos, self.lifetime

    def get_num_open(self, key, jqn=0, mqn=0, nmax=100, vmax=9999,
                     energy_offset=0.0):
        """
        simple wrapper around _read_data and _num_open
        Variables
        jqn       - total jqn quantum number.
        mqn     - mj_qn projection quantum number.
        nmax    - maximum allowed rotaional quantum number of the molecule.
        vmax    - maximum allowed vibrational quantum number of the molecule.
        energy_offset - calculate dos around energy = energy_offset
        """
        lmax = nmax + jqn
        self.bound_states_d, self.nmax = _read_data(
            self.systems[key]['dimer_dirn'],
            self.systems[key]['dimer_filen'],
            nmax)
        self.bound_states_c, self.lmax = _read_data(
            self.systems[key]['cmplx_dirn'],
            self.systems[key]['cmplx_filen'],
            lmax)
        self.num_open = AMDOS._num_open(
            jqn,
            mqn,
            self.bound_states_c,
            self.bound_states_d,
            vmax,
            energy_offset)
        return self.num_open

class MMDOS(object):
    """
    class for molecule molecule density of states calculations
    """
    def __init__(self, filen):
        self.systems = get_data(filen)
        self.nmax = None
        self.lmax = None
        self.bound_states_c = None
        self.bound_states_d = None
        self.lifetime = None
        self.dos = None

    @staticmethod
    def _mm_dos(jqn, mj_qn, nmax, vmax, lmax,
                bound_states_d, bound_states_c,
                energy_range=10.0):
        """
        computes the dos for an atom moleucle collsiions as described in
        Scattering of Ultracold Molecules in the Highly Resonant Regime --
        Mayle, Ruzic, Quememer and Bohn
        Variables
        jqn       - total J quantum number.
        mqn     - mj_qn projection quantum number.
        bound_states_c - boundstates of the complex obtained from _read_data
        bound_states_d - boundstates of the molecule obtained from _read_data
        vmax    - maximum vibrational quantum number.
        lmax    - maximum end-over-end rotational quantum number of the two
                  molecules
        nmax    - maximum rotational quantum number of a single molecule.
        energy_range  - dos = number of states counted/energy range
        """
        if abs(mj_qn) > abs(jqn):
            raise Exception('physically impossible (abs(mj_qn) > abs(jqn))')
        num = 0  # variable to hold the number of states between limits
        limit = (energy_range/2.0)*K2EH  # convert energy_range K to hartrees
        abs_gs = bound_states_d[0][0]  # energy of the absolute ground state
        for nqn in xrange(max(0, jqn-lmax), min(2*nmax, lmax-jqn)+1):
            for lqn in xrange(abs(jqn-nqn), min(lmax, jqn+nqn)+1):
                for n1qn in xrange(max(0, nqn-nmax), nmax+1):
                    for n2qn in xrange(abs(nqn-n1qn), min(nqn+n1qn, nmax)+1):
                        # looping over all vibrational levels of dimer1
                        for v1qn in xrange(min(len(bound_states_d[n1qn]),
                                               vmax+1)):
                            # looping over all vibrational levels of dimer2
                            for v2qn in xrange(min(len(bound_states_d[n2qn]),
                                                   vmax+1)):
                                threshold_energy = (
                                    bound_states_d[n1qn][v1qn] - abs_gs +
                                    bound_states_d[n2qn][v2qn] - abs_gs)
                                if (bound_states_c[lqn][0] >
                                        limit-threshold_energy):
                                    # the lowest state is higher than highest
                                    # threshold
                                    break
                                else:
                                    start_end = np.searchsorted(
                                        bound_states_c[lqn],
                                        [-limit-threshold_energy,
                                         limit-threshold_energy],
                                        'left')
                                    deg = (abs(max(-lqn, mj_qn-nqn)) +
                                           min(lqn, mj_qn+nqn) + 1)
                                    num += deg*(start_end[1] - start_end[0])
        # return dos in per uK i.e.times 10**-6K
        dos = (float(num)/energy_range)*1.0E-6
        # return lifetime in ms
        lifetime = dos*1.0E6*EH2K*2.0*pi*ATOMICUNITOFTIME * 1.0e3
        return dos, lifetime

    def mm_dos(self, key, jqn=0, mqn=0, nmax=100, vmax=999):
        """
        simple wrapper around _read_data and MMDOS
        Variables
        jqn       - total J quantum number.
        mqn     - mj_qn projection quantum number.
        nmax    - maximum allowed rotaional quantum number of the molecule.
        vmax    - maximum allowed vibrational quantum number of the molecule.
        energy_offset - calculate dos around energy = energy_offset
        """
        lmax = 2*nmax+jqn
        self.bound_states_d, self.nmax = _read_data(
            self.systems[key]['dimer_dirn'],
            self.systems[key]['dimer_filen'],
            nmax)
        self.bound_states_c, self.lmax = _read_data(
            self.systems[key]['cmplx_dirn'],
            self.systems[key]['cmplx_filen'],
            lmax)
        self.dos, self.lifetime = MMDOS._mm_dos(
            jqn, mqn, self.nmax, vmax, self.lmax,
            self.bound_states_d, self.bound_states_c)
        return self.dos, self.lifetime
<|endoftext|>


<commit_before>#!/usr/bin/python -tt

# omega - python client
# https://github.com/jfillmore/Omega-API-Engine
#
# Copyright 2011, Jonathon Fillmore
# Licensed under the MIT license. See LICENSE file.
# http://www.opensource.org/licenses/mit-license.php

"""Omega client for talking to an Omega server."""

import pycurl
import urllib
import json
import re
import dbg
import StringIO
import os
import tempfile
from error import Error

class OmegaClient:
	"""Client for talking to an Omega Server."""
	_version = '0.2'
	_request_type = 'POST'
	_curl = None
	_hostname = None
	_folder = '/'
	_url = None
	_session_coookie = None
	_cookie_file = None
	_useragent = 'OmegaClient/0.2'

	def __init__(self, url = 'localhost', credentials = None, port = 5800, use_https = True):
		self._cookie_file = os.path.expanduser('~/.omega_cookie') # tempfile.NamedTemporaryFile()
		self._curl = pycurl.Curl()
		self.set_url(url)
		self.set_credentials(credentials)
		self.set_port(port)
		self.set_https(use_https)
		# TODO: python 2.7 supports an order tuple object we can use to preserve order :)
		self.encode = json.JSONEncoder().encode
		self.decode = json.JSONDecoder().decode
		# setup cookie jar

	def set_url(self, url):
		if url != '':
			self._url = url
			if url.find('/') == -1:
				self._hostname = url
				self._folder = ''
			else:
				# check for protocol
				if url.lower()[0:7] == 'http://':
					url = url[8:]
					self.set_https(false)
				elif url.lower()[0:8] == 'https://':
					url = url[9:]
					self.set_https(true)
				else:
					# some other protocol perhaps?
					if re.match('^\w+://', url):
						raise Exception('Only HTTP and HTTPS are supported protocols.')
				# do we have a port?
				match = re.match('^(\w+\.)*\w+(:(\d{0,5}))/?', url)
				if match:
					self.set_port(match.group(3))
					# chop out the port
					url = '/'.join((re.split(':\d+/?', url)))
				(self._hostname, self._folder) = url.split('/', 1)
				# force ourself to end in a slash
				if not self._folder.endswith('/'):
					self._folder = ''.join((self._folder, '/'))
		else:
			raise Exception('Invalid API service URL: %s.' % url)

	def get_url(self):
		return self._url

	def get_hostname(self):
		return self._hostname

	def get_folder(self):
		return self._folder

	def get_request_type(self):
		return self._request_type;

	def set_request_type(self, type):
		if type.upper() in ('POST', 'GET'):
			self._request_type = type
		else:
			raise Exception('The request must be either "GET" or "POST".')

	def set_credentials(self, creds):
		if type(creds).__name__ == 'dict':
			# make sure we have a user/pass
			if 'username' in creds and 'password' in creds:
				self._credentials = creds
			elif 'token' in creds:
				self._credentials = creds
			else:
				raise Exception("Invalid credentials. Keys of 'username'/'password' or 'token' expected, but were not found.")
		elif creds == None:
			self._credentials = creds
		else:
			raise Exception("Invalid credentials. Keys of 'username'/'password' or 'token' expected, but were not found.")

	def set_https(self, secure = True):
		if secure:
			self._use_https = True
		else:
			self._use_https = False

	def set_port(self, port):
		if port >= 0 and port <= 65535:
			self._port = port
		else:
			raise Exception('Invalid API service port: %i.' % port)

	def run(self, api, args = (), raw_response = False, full_response = False, get = None, post = None, files = None):
		# check and prep the data
		if api == '':
			raise Exception("Invalid service API: '%s'." %api)
		api = urllib.quote(api)
		curl = pycurl.Curl()
		data = [
			('OMEGA_ENCODING', (curl.FORM_CONTENTS, 'json')),
			('OMEGA_API_PARAMS', (curl.FORM_CONTENTS, self.encode(args)))
		]
		if self._credentials:
			data.append(('OMEGA_CREDENTIALS', (curl.FORM_CONTENTS, self.encode(self._credentials))))
		# include any extra post data
		if post:
			(name, value) = post.split('=', 1)
			data.append((name, (curl.FORM_CONTENTS, value)))
		if files:
			# add in our files to the data
			for name in files:
				data.append((name, (curl.FORM_FILE, files[name])))
		# figure our our URL and get args
		url = self._url
		if self._use_https:
			url = ''.join(('https://', self._hostname))
		else:
			url = ''.join(('http://', self._hostname))
		url = '/'.join((':'.join((url, str(self._port))), self._folder))
		url = '/'.join((url, api))
		if get:
			url = '?'.join((url, get))
		# fire away
		curl.setopt(curl.URL, url)
		curl.setopt(curl.POST, 1)
		curl.setopt(curl.USERAGENT, self._useragent)
		curl.setopt(curl.COOKIEFILE, self._cookie_file)
		curl.setopt(curl.COOKIEJAR, self._cookie_file)
		if self._use_https:
			curl.setopt(curl.SSL_VERIFYPEER, 0) # TODO: don't always assume
			curl.setopt(curl.SSL_VERIFYHOST, 0) # TODO: don't always assume
		if data:
			curl.setopt(curl.HTTPPOST, data)
		else:
			curl.setopt(curl.POSTFIELDS, '&'.join(args))
		response = StringIO.StringIO()
		curl.setopt(curl.WRITEFUNCTION, response.write)
		curl.perform()
		response = response.getvalue()
		http_code = curl.getinfo(curl.HTTP_CODE)
		if http_code != 200:
			raise Exception("Failed to execute API %s: server returned HTTP code %s" % (api, str(http_code)))
		curl.close()
		if raw_response:
			return response
		else:
			# decode the response and check whether or not it was successful
			# TODO: check response encoding in header
			try:
				response = self.decode(response)
			except:
				raise Error('Failed to decode API response.', response)
			# check to see if our API call was successful
			if 'result' in response and response['result'] == False:
				if 'reason' in response:
					if full_response:
						raise Error('API "%s" failed.\n%s' % (urllib.unquote(api), dbg.obj2str(response)))
					else:
						raise Error(response['reason'])
				else:
					raise Error('API "%s" failed, but did not provide an explanation. Response: %s' % (api, response))
			else:
				if full_response:
					return response
				else:
					if 'data' in response:
						return response['data']
					else:
						return None

if __name__ == '__main__':
	import dbg
	dbg.pretty_print(OmegaClient());
<commit_msg>allow simplejson fallback if json module not available<commit_after>#!/usr/bin/python -tt

# omega - python client
# https://github.com/jfillmore/Omega-API-Engine
#
# Copyright 2011, Jonathon Fillmore
# Licensed under the MIT license. See LICENSE file.
# http://www.opensource.org/licenses/mit-license.php

"""Omega client for talking to an Omega server."""

import pycurl
import urllib
try:
    import json
except:
    import simplejson
    json = simplejson
import re
import dbg
import StringIO
import os
import tempfile
from error import Error

class OmegaClient:
	"""Client for talking to an Omega Server."""
	_version = '0.2'
	_request_type = 'POST'
	_curl = None
	_hostname = None
	_folder = '/'
	_url = None
	_session_coookie = None
	_cookie_file = None
	_useragent = 'OmegaClient/0.2'

	def __init__(self, url = 'localhost', credentials = None, port = 5800, use_https = True):
		self._cookie_file = os.path.expanduser('~/.omega_cookie') # tempfile.NamedTemporaryFile()
		self._curl = pycurl.Curl()
		self.set_url(url)
		self.set_credentials(credentials)
		self.set_port(port)
		self.set_https(use_https)
		# TODO: python 2.7 supports an order tuple object we can use to preserve order :)
		self.encode = json.JSONEncoder().encode
		self.decode = json.JSONDecoder().decode
		# setup cookie jar

	def set_url(self, url):
		if url != '':
			self._url = url
			if url.find('/') == -1:
				self._hostname = url
				self._folder = ''
			else:
				# check for protocol
				if url.lower()[0:7] == 'http://':
					url = url[8:]
					self.set_https(false)
				elif url.lower()[0:8] == 'https://':
					url = url[9:]
					self.set_https(true)
				else:
					# some other protocol perhaps?
					if re.match('^\w+://', url):
						raise Exception('Only HTTP and HTTPS are supported protocols.')
				# do we have a port?
				match = re.match('^(\w+\.)*\w+(:(\d{0,5}))/?', url)
				if match:
					self.set_port(match.group(3))
					# chop out the port
					url = '/'.join((re.split(':\d+/?', url)))
				(self._hostname, self._folder) = url.split('/', 1)
				# force ourself to end in a slash
				if not self._folder.endswith('/'):
					self._folder = ''.join((self._folder, '/'))
		else:
			raise Exception('Invalid API service URL: %s.' % url)

	def get_url(self):
		return self._url

	def get_hostname(self):
		return self._hostname

	def get_folder(self):
		return self._folder

	def get_request_type(self):
		return self._request_type;

	def set_request_type(self, type):
		if type.upper() in ('POST', 'GET'):
			self._request_type = type
		else:
			raise Exception('The request must be either "GET" or "POST".')

	def set_credentials(self, creds):
		if type(creds).__name__ == 'dict':
			# make sure we have a user/pass
			if 'username' in creds and 'password' in creds:
				self._credentials = creds
			elif 'token' in creds:
				self._credentials = creds
			else:
				raise Exception("Invalid credentials. Keys of 'username'/'password' or 'token' expected, but were not found.")
		elif creds == None:
			self._credentials = creds
		else:
			raise Exception("Invalid credentials. Keys of 'username'/'password' or 'token' expected, but were not found.")

	def set_https(self, secure = True):
		if secure:
			self._use_https = True
		else:
			self._use_https = False

	def set_port(self, port):
		if port >= 0 and port <= 65535:
			self._port = port
		else:
			raise Exception('Invalid API service port: %i.' % port)

	def run(self, api, args = (), raw_response = False, full_response = False, get = None, post = None, files = None):
		# check and prep the data
		if api == '':
			raise Exception("Invalid service API: '%s'." %api)
		api = urllib.quote(api)
		curl = pycurl.Curl()
		data = [
			('OMEGA_ENCODING', (curl.FORM_CONTENTS, 'json')),
			('OMEGA_API_PARAMS', (curl.FORM_CONTENTS, self.encode(args)))
		]
		if self._credentials:
			data.append(('OMEGA_CREDENTIALS', (curl.FORM_CONTENTS, self.encode(self._credentials))))
		# include any extra post data
		if post:
			(name, value) = post.split('=', 1)
			data.append((name, (curl.FORM_CONTENTS, value)))
		if files:
			# add in our files to the data
			for name in files:
				data.append((name, (curl.FORM_FILE, files[name])))
		# figure our our URL and get args
		url = self._url
		if self._use_https:
			url = ''.join(('https://', self._hostname))
		else:
			url = ''.join(('http://', self._hostname))
		url = '/'.join((':'.join((url, str(self._port))), self._folder))
		url = '/'.join((url, api))
		if get:
			url = '?'.join((url, get))
		# fire away
		curl.setopt(curl.URL, url)
		curl.setopt(curl.POST, 1)
		curl.setopt(curl.USERAGENT, self._useragent)
		curl.setopt(curl.COOKIEFILE, self._cookie_file)
		curl.setopt(curl.COOKIEJAR, self._cookie_file)
		if self._use_https:
			curl.setopt(curl.SSL_VERIFYPEER, 0) # TODO: don't always assume
			curl.setopt(curl.SSL_VERIFYHOST, 0) # TODO: don't always assume
		if data:
			curl.setopt(curl.HTTPPOST, data)
		else:
			curl.setopt(curl.POSTFIELDS, '&'.join(args))
		response = StringIO.StringIO()
		curl.setopt(curl.WRITEFUNCTION, response.write)
		curl.perform()
		response = response.getvalue()
		http_code = curl.getinfo(curl.HTTP_CODE)
		if http_code != 200:
			raise Exception("Failed to execute API %s: server returned HTTP code %s" % (api, str(http_code)))
		curl.close()
		if raw_response:
			return response
		else:
			# decode the response and check whether or not it was successful
			# TODO: check response encoding in header
			try:
				response = self.decode(response)
			except:
				raise Error('Failed to decode API response.', response)
			# check to see if our API call was successful
			if 'result' in response and response['result'] == False:
				if 'reason' in response:
					if full_response:
						raise Error('API "%s" failed.\n%s' % (urllib.unquote(api), dbg.obj2str(response)))
					else:
						raise Error(response['reason'])
				else:
					raise Error('API "%s" failed, but did not provide an explanation. Response: %s' % (api, response))
			else:
				if full_response:
					return response
				else:
					if 'data' in response:
						return response['data']
					else:
						return None

if __name__ == '__main__':
	import dbg
	dbg.pretty_print(OmegaClient());
<|endoftext|>


<commit_before>import random

class Dice(object):
    def __init__(self):
        self.score = 0
        self.set_aside = []
        self.remaining = [0, 0, 0, 0, 0, 0]

    def get_score(self):
        return self.score

    def get_set_aside(self, as_str=False):
        if as_str:
            return Dice.to_str(self.set_aside)
        else:
            return tuple(self.set_aside)

    def get_remaining(self, as_str=False):
        if as_str:
            return Dice.to_str(self.remaining)
        else:
            return tuple(self.remaining)

    @staticmethod
    def to_str(dice_values):
        return ' '.join([str(die) for die in dice_values])

    def roll(self):
        self.check_for_farkle(self.remaining) #<- eh? does not appear to do anything
        self.remaining = [random.randint(1,6) for die in self.remaining]

    #DOES NOT WORK for all cases
    def is_valid_set_aside(self, dice_values):
        for die_value in dice_values:
            found = False
            for die in self.remaining:
                if die == die_value:
                    found = True
            if not found:
                return False
        return True

    @staticmethod
    def find_n_of_a_kind(n, die_counts):
        matches = []
        for i in range(1,7):
            if die_counts[i] >= n:
                matches.append(i)
        return tuple(matches)

    #@staticmethod
    #def zero_out_counts(die_counts):   <- ??

    def dice_combination_value(self, dice_values):
        score = 0
        die_counts = {1:0, 2:0, 3:0, 4:0, 5:0, 6:0}

        for die in dice_values:
            die_counts[die] += 1

        # four with a pair
        matches4, matches2 = Dice.find_n_of_a_kind(4, die_counts), Dice.find_n_of_a_kind(2, die_counts)
        if matches4 and matches2:
            score += 1500
            die_counts[matches4[0]] = die_counts[matches2[0]] = 0

        # two triplets
        matches = Dice.find_n_of_a_kind(3, die_counts)
        if len(matches) == 2
            score += 2500
            die_counts[matches[0]] = die_counts[matches[1]] = 0

        # three pairs
        matches = Dice.find_n_of_a_kind(2, die_counts)
        if len(matches) == 3:
            score += 1500
            die_counts[matches[0]] = die_counts[matches[1]] = die_counts[matches[2]] = 0

        # strait
        if len(set(range(1,7)) - set(dice_values)) == 0:
            score += 1500
            for i in range(1,7):
                die_counts[i] = 0

        # 6, 5, 4, and 3 of a kind
        matches = Dice.find_n_of_a_kind(6, die_counts)
        if matches:
            score += 3000
            die_counts[matches[0]] = 0

        matches = Dice.find_n_of_a_kind(5, die_counts)
        if matches:
            score += 2000
            die_counts[matches[0]] = 0

        matches = Dice.find_n_of_a_kind(4, die_counts)
        if matches:
            score += 1000
            die_counts[matches[0]] = 0

        matches = Dice.find_n_of_a_kind(3, die_counts)
        if matches[0] == 1:
            score += 300
        else:
            score += matches[0] * 100
        die_counts[matches[0]] = 0

        # single 1's and 5's
        if die_counts[1] > 0:
            score += die_counts[1] * 100
            die_counts[1] = 0
        if die_counts[5] > 0:
            score += die_counts[5] * 50
            die_counts[5] = 0

#        rules = [
#                {'die_counts':(4,2),'value':1500},
#                {'die_counts':(3,3),'value':2500},
#                {'die_counts':(2,2,2),'value':1500},
#                {'die_counts':(6,),'value':3000},
#                {'die_counts':(5,),'value':2000},
#                {'die_counts':(4,),'value':1000},
#                {'die_counts':(
#
#                {'die_counts':(3,),'value':},
#abstract out the setting of the die_counts to 0 ?

        if any(die_counts.values()):
            raise InvalidSetAsideException()

        return (score, die_counts)
        #return score and also remaining dice, do exception checking in two separate calling functions

    def check_for_farkle(self, dice_values):
        value, leftovers = self.dice_combination_value(dice_values)
        if value = 0:
            raise GotFarkleException()

    def evaluate_set_aside(self, dice_values):
        value, leftovers = self.dice_combination_value(dice_values)
        if value > 0 and not any(leftovers.values()):
            raise InvalidSetAsideException()
        else:
            return value

    def move_to_set_aside(self, dice_values):

        value = self.evaluate_set_aside(dice_values)
        self.score += value

        for die_value in dice_values:
            self.set_aside.append(die_value)
            self.remaining.remove(die_value)

        if len(self.set_aside) == 6:
            self.remaining = self.set_aside
            self.set_aside = []

class Rules(object):

class GotFarkleException(Exception):
    pass

class InvalidSetAsideException(Exception):
    pass

class CantRollException(Exception):
    pass

class HumanPlayer(object):
    def take_turn(self, dice, scores):
        while True:
            print "\n"*64
            print "\n\nScores:\n"
            for i, score in enumerate(scores):
                print "Player {0}: {1}".format(i, score)

            print "Turn score: ", dice.get_score()

            print "\nYou roll the dice:"
            try:
                dice.roll()
            except GotFarkleException as e:
                print "\nYou got a farkle!"
#HERE working on getting the interface to recognize a farkle, still print out the dice, then end the turn, also do not allow the dice to be rolled again once this happens
            print dice.get_remaining(as_str=True)

            print "\nSet Aside:"
            print dice.get_set_aside(as_str=True)

            while True:
                choices = raw_input("\nIndicate the dice you want to set aside by entering their numbers separated by spaces, or enter nothing to stop.\n")

                if choices == '':
                    return dice

                try:
                    dice.move_to_set_aside([int(choice) for choice in choices.split()])
                    break
                except ValueError as e:
                    print "The set aside must contain only integers from 1-6."
                except InvalidSetAsideException as e:
                    print "That set aside is not valid."

class AIPlayer(object):
    def take_turn(self, dice):
        pass

class Farkle(object):
    def __init__(self):
        self.players = []
        self.scores = []

    def add_player(self, player):
        self.players.append(player)
        self.scores.append(0)

#add a name, use the uuid for nonhuman players, could be a database key
    def play(self):

        turn_index = 0
        while True:
            dice = self.players[turn_index].take_turn(Dice(), tuple(self.scores))
            self.scores[turn_index] += dice.get_score()

            if self.scores[turn_index] > 10000: break
            turn_index += 1
            if turn_index > len(self.players) - 1: turn_index = 0

        return turn_index

def main():
    farkle = Farkle()
    farkle.add_player(HumanPlayer())
    winner = farkle.play()
    print "The winner is player " + str(winner) + "!"

if __name__ == "__main__":
    main()
<commit_msg>Fixed dice_in_set_aside function<commit_after>import random

class Dice(object):
    def __init__(self):
        self.score = 0
        self.set_aside = []
        self.remaining = [0, 0, 0, 0, 0, 0]
        self.got_farkle = False

    def get_score(self):
        return self.score

    def get_set_aside(self, as_str=False):
        if as_str:
            return Dice.to_str(self.set_aside)
        else:
            return tuple(self.set_aside)

    def get_remaining(self, as_str=False):
        if as_str:
            return Dice.to_str(self.remaining)
        else:
            return tuple(self.remaining)

    @staticmethod
    def to_str(dice_values):
        return ' '.join([str(die) for die in dice_values])

    def roll(self):
        if self.got_farkle:
            raise GotFarkleException()

        self.remaining = [random.randint(1,6) for die in self.remaining]
        self.check_for_farkle(self.remaining)

    #DOES NOT WORK for all cases
    def dice_in_set_aside(self, dice_values):

        proposed_dice = list(dice_values)
        actual_dice = list(self.remaining)

        remaining = list(dice_values)

        for die in actual_dice:
            if die in proposed_dice:
                proposed_dice.remove(die)
                actual_dice.remove(die)
        if len(proposed_dice) > 0:
            return False
        else:
            return True

    @staticmethod
    def find_n_of_a_kind(n, die_counts):
        matches = []
        for i in range(1,7):
            if die_counts[i] >= n:
                matches.append(i)
        return tuple(matches)

    #@staticmethod
    #def zero_out_counts(die_counts):   <- ??

    def dice_combination_value(self, dice_values):
        score = 0
        die_counts = {1:0, 2:0, 3:0, 4:0, 5:0, 6:0}

        for die in dice_values:
            die_counts[die] += 1

        # four with a pair
        matches4, matches2 = Dice.find_n_of_a_kind(4, die_counts), Dice.find_n_of_a_kind(2, die_counts)
        if matches4 and matches2:
            score += 1500
            die_counts[matches4[0]] = die_counts[matches2[0]] = 0

        # two triplets
        matches = Dice.find_n_of_a_kind(3, die_counts)
        if len(matches) == 2:
            score += 2500
            die_counts[matches[0]] = die_counts[matches[1]] = 0

        # three pairs
        matches = Dice.find_n_of_a_kind(2, die_counts)
        if len(matches) == 3:
            score += 1500
            die_counts[matches[0]] = die_counts[matches[1]] = die_counts[matches[2]] = 0

        # strait
        if len(set(range(1,7)) - set(dice_values)) == 0:
            score += 1500
            for i in range(1,7):
                die_counts[i] = 0

        # 6, 5, 4, and 3 of a kind
        matches = Dice.find_n_of_a_kind(6, die_counts)
        if matches:
            score += 3000
            die_counts[matches[0]] = 0

        matches = Dice.find_n_of_a_kind(5, die_counts)
        if matches:
            score += 2000
            die_counts[matches[0]] = 0

        matches = Dice.find_n_of_a_kind(4, die_counts)
        if matches:
            score += 1000
            die_counts[matches[0]] = 0

        matches = Dice.find_n_of_a_kind(3, die_counts)
        if matches:
            if matches[0] == 1:
                score += 300
            else:
                score += matches[0] * 100
            die_counts[matches[0]] = 0

        # single 1's and 5's
        if die_counts[1] > 0:
            score += die_counts[1] * 100
            die_counts[1] = 0
        if die_counts[5] > 0:
            score += die_counts[5] * 50
            die_counts[5] = 0

        return (score, die_counts)

#        rules = [
#                {'die_counts':(4,2),'value':1500},
#                {'die_counts':(3,3),'value':2500},
#                {'die_counts':(2,2,2),'value':1500},
#                {'die_counts':(6,),'value':3000},
#                {'die_counts':(5,),'value':2000},
#                {'die_counts':(4,),'value':1000},
#                {'die_counts':(
#
#                {'die_counts':(3,),'value':},
#abstract out the setting of the die_counts to 0 ?

    def check_for_farkle(self, dice_values):
        if dice_values == [0,0,0,0,0,0]: return
        value, leftovers = self.dice_combination_value(dice_values)
        if value == 0:
            self.score = 0
            self.got_farkle = True   # set so that the dice can't be rolled again after getting a farkle
            raise GotFarkleException()

    def evaluate_set_aside(self, dice_values):
        value, leftovers = self.dice_combination_value(dice_values)
        if value > 0 and not any(leftovers.values()):
            return value
        else:
            raise InvalidSetAsideException()

    def move_to_set_aside(self, dice_values):
        if not self.dice_in_set_aside(dice_values):
            raise InvalidSetAsideException()

        value = self.evaluate_set_aside(dice_values)
        self.score += value

        for die_value in dice_values:
            self.set_aside.append(die_value)
            self.remaining.remove(die_value)

        if len(self.set_aside) == 6:
            self.remaining = self.set_aside
            self.set_aside = []

class GotFarkleException(Exception):
    pass

class InvalidSetAsideException(Exception):
    pass

class CantRollException(Exception):
    pass

class HumanPlayer(object):
    def take_turn(self, dice, scores):
        while True:
            print "\n\nScores:\n"
            for i, score in enumerate(scores):
                print "Player {0}: {1}".format(i, score)

            print "Turn score: ", dice.get_score()

            print "\nSet Aside:"
            print dice.get_set_aside(as_str=True)

            print "\nYou roll the dice:"
            try:
                dice.roll()
            except GotFarkleException as e:
                print dice.get_remaining(as_str=True)
                print "\nYou got a farkle!"
                raw_input('Hit enter to end your turn.')
                return dice

            print dice.get_remaining(as_str=True)

            while True:
                choices = raw_input("\nIndicate the dice you want to set aside by entering their numbers separated by spaces, or enter nothing to stop.\n")

                try:
                    dice.move_to_set_aside([int(choice) for choice in choices.split()])
                    break
                except ValueError as e:
                    print "The set aside must contain only integers from 1-6."
                except InvalidSetAsideException as e:
                    print "That set aside is not valid."

            while True:
                choice = raw_input("You have {0} points.  Hit enter to continue rolling, or type 'stop' to end your turn.\n".format(dice.get_score()))
                if choice == '':
                    break
                if choice.lower() == 'stop':
                    return dice

class AIPlayer(object):
    def take_turn(self, dice):
        pass

class Farkle(object):
    def __init__(self):
        self.players = []
        self.scores = []

    def add_player(self, player):
        self.players.append(player)
        self.scores.append(0)

#add a name, use the uuid for nonhuman players, could be a database key
    def play(self):

        turn_index = 0
        while True:
            dice = self.players[turn_index].take_turn(Dice(), tuple(self.scores))
            self.scores[turn_index] += dice.get_score()

            if self.scores[turn_index] > 10000: break
            turn_index += 1
            if turn_index > len(self.players) - 1: turn_index = 0

        return turn_index

def main():
    farkle = Farkle()
    farkle.add_player(HumanPlayer())
    winner = farkle.play()
    print "The winner is player {0}!".format(winner)

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


<commit_before>#!/usr/bin/python
#
# Prototype for the beng-proxy translation server.
#
# Author: Max Kellermann <mk@cm4all.com>

from twisted.internet import reactor
from twisted.internet.protocol import Protocol, Factory
import struct

TRANSLATE_BEGIN = 1
TRANSLATE_END = 2
TRANSLATE_HOST = 3
TRANSLATE_URI = 4
TRANSLATE_STATUS = 5
TRANSLATE_PATH = 6
TRANSLATE_CONTENT_TYPE = 7
TRANSLATE_PROXY = 8
TRANSLATE_REDIRECT = 9
TRANSLATE_FILTER = 10
TRANSLATE_PROCESS = 11
TRANSLATE_SESSION = 12
TRANSLATE_PARAM = 13
TRANSLATE_USER = 14
TRANSLATE_LANGUAGE = 15
TRANSLATE_REMOTE_HOST = 16
TRANSLATE_PATH_INFO = 17
TRANSLATE_SITE = 18

class PacketReader:
    def __init__(self):
        self._header = ''
        self.complete = False

    def consume(self, data):
        assert not self.complete
        assert isinstance(data, str)
        # read header
        if len(self._header) < 4:
            # append to header
            size = 4 - len(self._header)
            self._header += data[0:size]
            data = data[size:]
            if len(self._header) < 4:
                return data

            # header is finished, decode it
            self._length, self.command = struct.unpack('HH', self._header)
            self.payload = ''
            if self._length == 0:
                # no payload, we're done
                self.complete = True
                return data

        # read payload
        if len(self.payload) < self._length:
            # append to payload
            size = self._length - len(self.payload)
            self.payload += data[0:size]
            data = data[size:]

            # done yet?
            if len(self.payload) == self._length:
                self.complete = True

        # return data chunk without the consumed part
        return data

class Request:
    def __init__(self):
        self.host = None
        self.uri = None
        self.session = None
        self.param = None
        self.remote_host = None

    def packetReceived(self, packet):
        if packet.command == TRANSLATE_END:
            return True
        elif packet.command == TRANSLATE_HOST:
            self.host = packet.payload
        elif packet.command == TRANSLATE_URI:
            self.uri = packet.payload
        elif packet.command == TRANSLATE_SESSION:
            self.session = packet.payload
        elif packet.command == TRANSLATE_PARAM:
            self.param = packet.payload
        elif packet.command == TRANSLATE_REMOTE_HOST:
            self.remote_host = packet.payload
        else:
            print "Invalid command:", packet.command
        return False

class Translation(Protocol):
    def connectionMade(self):
        print "Connected from", self.transport.client
        self._request = None
        self._packet = None

    def _write_packet(self, command, payload = ''):
        assert isinstance(payload, str)
        self.transport.write(struct.pack('HH', len(payload), command))
        self.transport.write(payload)

    def _handle_request(self, request):
        if request.session is not None: print "- session =", request.session
        if request.param is not None: print "- param =", request.param

        if request.param is not None:
            # log in or log out; "real" authentification is missing
            # here.  We're logging out if request.param is an empty
            # string.
            user = session = request.param
        elif request.session is not None:
            # user is already authenticated
            user = request.session
            session = None
        else:
            #
            user = session = None

        if request.uri == '/beng-proxy.js':
            from sys import argv
            from os.path import abspath, dirname, join
            path = join(dirname(dirname(abspath(argv[0]))), 'js/beng-proxy.js')
            print path
        else:
            path = '/var/www' + request.uri

        self._write_packet(TRANSLATE_BEGIN)
        self._write_packet(TRANSLATE_PATH, path)
        if user is not None:
            self._write_packet(TRANSLATE_USER, user)
        if session is not None:
            self._write_packet(TRANSLATE_SESSION, session)
        self._write_packet(TRANSLATE_FILTER, 'http://cfatest01.intern.cm-ag/filter.py')
        self._write_packet(TRANSLATE_PROCESS)
        self._write_packet(TRANSLATE_END)

    def _handle_packet(self, packet):
        if packet.command == TRANSLATE_BEGIN:
            self._request = Request()
        elif self._request is not None:
            if self._request.packetReceived(packet):
                self._handle_request(self._request)
                self._request = None
        else:
            print "Invalid command without request:", packet.command

    def dataReceived(self, data):
        while len(data) > 0:
            if self._packet is None:
                self._packet = PacketReader()
            data = self._packet.consume(data)
            assert len(data) == 0 or self._packet.complete
            if self._packet.complete:
                self._handle_packet(self._packet)
                self._packet = None

    def connectionLost(self, reason):
        print "Disconnected from", self.transport.client

factory = Factory()
factory.protocol = Translation

if __name__ == '__main__':
    reactor.listenUNIX('/tmp/beng-proxy-translate', factory)
    reactor.run()
<commit_msg>disable filter for now<commit_after>#!/usr/bin/python
#
# Prototype for the beng-proxy translation server.
#
# Author: Max Kellermann <mk@cm4all.com>

from twisted.internet import reactor
from twisted.internet.protocol import Protocol, Factory
import struct

TRANSLATE_BEGIN = 1
TRANSLATE_END = 2
TRANSLATE_HOST = 3
TRANSLATE_URI = 4
TRANSLATE_STATUS = 5
TRANSLATE_PATH = 6
TRANSLATE_CONTENT_TYPE = 7
TRANSLATE_PROXY = 8
TRANSLATE_REDIRECT = 9
TRANSLATE_FILTER = 10
TRANSLATE_PROCESS = 11
TRANSLATE_SESSION = 12
TRANSLATE_PARAM = 13
TRANSLATE_USER = 14
TRANSLATE_LANGUAGE = 15
TRANSLATE_REMOTE_HOST = 16
TRANSLATE_PATH_INFO = 17
TRANSLATE_SITE = 18

class PacketReader:
    def __init__(self):
        self._header = ''
        self.complete = False

    def consume(self, data):
        assert not self.complete
        assert isinstance(data, str)
        # read header
        if len(self._header) < 4:
            # append to header
            size = 4 - len(self._header)
            self._header += data[0:size]
            data = data[size:]
            if len(self._header) < 4:
                return data

            # header is finished, decode it
            self._length, self.command = struct.unpack('HH', self._header)
            self.payload = ''
            if self._length == 0:
                # no payload, we're done
                self.complete = True
                return data

        # read payload
        if len(self.payload) < self._length:
            # append to payload
            size = self._length - len(self.payload)
            self.payload += data[0:size]
            data = data[size:]

            # done yet?
            if len(self.payload) == self._length:
                self.complete = True

        # return data chunk without the consumed part
        return data

class Request:
    def __init__(self):
        self.host = None
        self.uri = None
        self.session = None
        self.param = None
        self.remote_host = None

    def packetReceived(self, packet):
        if packet.command == TRANSLATE_END:
            return True
        elif packet.command == TRANSLATE_HOST:
            self.host = packet.payload
        elif packet.command == TRANSLATE_URI:
            self.uri = packet.payload
        elif packet.command == TRANSLATE_SESSION:
            self.session = packet.payload
        elif packet.command == TRANSLATE_PARAM:
            self.param = packet.payload
        elif packet.command == TRANSLATE_REMOTE_HOST:
            self.remote_host = packet.payload
        else:
            print "Invalid command:", packet.command
        return False

class Translation(Protocol):
    def connectionMade(self):
        print "Connected from", self.transport.client
        self._request = None
        self._packet = None

    def _write_packet(self, command, payload = ''):
        assert isinstance(payload, str)
        self.transport.write(struct.pack('HH', len(payload), command))
        self.transport.write(payload)

    def _handle_request(self, request):
        if request.session is not None: print "- session =", request.session
        if request.param is not None: print "- param =", request.param

        if request.param is not None:
            # log in or log out; "real" authentification is missing
            # here.  We're logging out if request.param is an empty
            # string.
            user = session = request.param
        elif request.session is not None:
            # user is already authenticated
            user = request.session
            session = None
        else:
            #
            user = session = None

        if request.uri == '/beng-proxy.js':
            from sys import argv
            from os.path import abspath, dirname, join
            path = join(dirname(dirname(abspath(argv[0]))), 'js/beng-proxy.js')
            print path
        else:
            path = '/var/www' + request.uri

        self._write_packet(TRANSLATE_BEGIN)
        self._write_packet(TRANSLATE_PATH, path)
        if user is not None:
            self._write_packet(TRANSLATE_USER, user)
        if session is not None:
            self._write_packet(TRANSLATE_SESSION, session)
        #self._write_packet(TRANSLATE_FILTER, 'http://cfatest01.intern.cm-ag/filter.py')
        self._write_packet(TRANSLATE_PROCESS)
        self._write_packet(TRANSLATE_END)

    def _handle_packet(self, packet):
        if packet.command == TRANSLATE_BEGIN:
            self._request = Request()
        elif self._request is not None:
            if self._request.packetReceived(packet):
                self._handle_request(self._request)
                self._request = None
        else:
            print "Invalid command without request:", packet.command

    def dataReceived(self, data):
        while len(data) > 0:
            if self._packet is None:
                self._packet = PacketReader()
            data = self._packet.consume(data)
            assert len(data) == 0 or self._packet.complete
            if self._packet.complete:
                self._handle_packet(self._packet)
                self._packet = None

    def connectionLost(self, reason):
        print "Disconnected from", self.transport.client

factory = Factory()
factory.protocol = Translation

if __name__ == '__main__':
    reactor.listenUNIX('/tmp/beng-proxy-translate', factory)
    reactor.run()
<|endoftext|>


<commit_before>import sys
sys.path.append('../../DeepLearningTutorials/code/')

import numpy
import theano
import theano.tensor as T

import random

def get_samples_with_target(source_num = 0, data_dim = 0, n_flag=0):

    if source_num is 0:
        [x, y] = normal_dist(data_dim, 100, 100, [2,8], seed = 1)

    elif source_num is 1:
        from sklearn import datasets
        iris = datasets.load_iris()
        x, y = chop_data(iris.data, iris.target, data_dim, n_flag)

    elif source_num is 2:
        from logistic_sgd import load_data
        datasets = load_data('mnist.pkl.gz')
        [shared_x, shared_y] = datasets[0]
        x, y = chop_data(shared_x.get_value(), shared_y.eval(), data_dim, n_flag)

    else:
        raise ValueError

    return x, y

def chop_data(all_data, all_target, data_dim, n_flag):
    """ reduce the number of category of the flags to n_flag """
    all_flag = numpy.unique(all_target)
    flags = all_flag[0: min(all_flag.shape[0], n_flag)]

    x_list = []
    y_list = []
    for flag in flags:
        x_list.append(all_data[all_target == flag])
        y_list.append(all_target[all_target == flag])

    x = numpy.vstack(x_list)
    y = numpy.hstack(y_list)

    if data_dim > 0:
        x = x[:, 0:data_dim]

    return x,y

def split_to_three_sets(x, y, valid_ratio = 1./3, test_ratio = 1./3, r_seed = 1):
    n_valid = int(valid_ratio * x.shape[0])
    n_test = int(test_ratio * x.shape[0])
    n_train = x.shape[0] - n_valid - n_test

    index_all = range( x.shape[0])
    random.seed( r_seed)
    random.shuffle( index_all)

    train_index = index_all[0:n_train]
    valid_index = index_all[n_train:n_train+n_valid]
    test_index = index_all[n_train+n_valid: n_train+n_valid+n_test]

    train_x = x[train_index, :]
    train_y = y[train_index]

    valid_x = x[valid_index, :]
    valid_y = y[valid_index]

    test_x = x[test_index, :]
    test_y = y[test_index]

    return [train_x, train_y, valid_x, valid_y, test_x, test_y]

def shared_array(set_x):
    shared_x = theano.shared(
        numpy.asarray(set_x, dtype = theano.config.floatX),
        borrow = True)
    return shared_x

def shared_flag(set_y):
    shared_y = theano.shared(
        numpy.asarray(set_y, dtype=theano.config.floatX),
        borrow=True)
    return T.cast(shared_y, 'int32')

def normal_dist(n_dim = 2, n_neg_sample = 100, n_pos_sample = 100, bias = [-2, 2], seed = 1):

    """ Generate 2 element samples of normal distribution """
    data = []

    random.seed(seed)
    numpy.random.seed(seed)

    for i in xrange(0,n_neg_sample):
        vec = numpy.random.randn(1,n_dim) + bias[0]
        flag = 0
        data.append([vec,flag])
    for i in range(0,n_pos_sample):
        vec = numpy.random.randn(1,n_dim) + bias[1]
        flag = 1
        data.append([vec,flag])

    random.shuffle(data)

    x = numpy.array([item[0][0] for item in data])
    y = numpy.array([item[1] for item in data])

    return [x,y]

def uniform_dist(n_dim = 2, n_sample = 100, minimum = 0.0, maximum = 1.0, seed = 1):

    numpy.random.seed(seed)
    return numpy.random.uniform(minimum, maximum, (n_sample, n_dim))
<commit_msg>Implement time-series generator<commit_after>import sys
sys.path.append('../../DeepLearningTutorials/code/')

import numpy as np
import theano
import theano.tensor as T

import random

def generate_time_series(source_num = 0):

    if source_num is 0:
        [x, y] = normal_timeseries()
    else:
        raise ValueError('source_num must be 0')
    return x, y

def get_samples_with_target(source_num = 0, data_dim = 0, n_flag=0):

    if source_num is 0:
        [x, y] = normal_dist(data_dim, 100, 100, [2,8], seed = 1)

    elif source_num is 1:
        from sklearn import datasets
        iris = datasets.load_iris()
        x, y = chop_data(iris.data, iris.target, data_dim, n_flag)

    elif source_num is 2:
        from logistic_sgd import load_data
        datasets = load_data('mnist.pkl.gz')
        [shared_x, shared_y] = datasets[0]
        x, y = chop_data(shared_x.get_value(), shared_y.eval(), data_dim, n_flag)

    else:
        raise ValueError

    return x, y

def chop_data(all_data, all_target, data_dim, n_flag):
    """ reduce the number of category of the flags to n_flag """
    all_flag = numpy.unique(all_target)
    flags = all_flag[0: min(all_flag.shape[0], n_flag)]

    x_list = []
    y_list = []
    for flag in flags:
        x_list.append(all_data[all_target == flag])
        y_list.append(all_target[all_target == flag])

    x = np.vstack(x_list)
    y = np.hstack(y_list)

    if data_dim > 0:
        x = x[:, 0:data_dim]

    return x,y

def split_to_three_sets(x, y, valid_ratio = 1./3, test_ratio = 1./3, r_seed = 1):
    n_valid = int(valid_ratio * x.shape[0])
    n_test = int(test_ratio * x.shape[0])
    n_train = x.shape[0] - n_valid - n_test

    index_all = range( x.shape[0])
    random.seed( r_seed)
    random.shuffle( index_all)

    train_index = index_all[0:n_train]
    valid_index = index_all[n_train:n_train+n_valid]
    test_index = index_all[n_train+n_valid: n_train+n_valid+n_test]

    train_x = x[train_index, :]
    train_y = y[train_index]

    valid_x = x[valid_index, :]
    valid_y = y[valid_index]

    test_x = x[test_index, :]
    test_y = y[test_index]

    return [train_x, train_y, valid_x, valid_y, test_x, test_y]

def shared_array(set_x):
    shared_x = theano.shared(
        np.asarray(set_x, dtype = theano.config.floatX),
        borrow = True)
    return shared_x

def shared_flag(set_y):
    shared_y = theano.shared(
        np.asarray(set_y, dtype=theano.config.floatX),
        borrow=True)
    return T.cast(shared_y, 'int32')

def normal_dist(n_dim = 2, n_neg_sample = 100, n_pos_sample = 100, bias = [-2, 2], seed = 1):

    """ Generate 2 element samples of normal distribution """
    data = []

    random.seed(seed)
    np.random.seed(seed)

    for i in xrange(0,n_neg_sample):
        vec = np.random.randn(1,n_dim) + bias[0]
        flag = 0
        data.append([vec,flag])
    for i in range(0,n_pos_sample):
        vec = np.random.randn(1,n_dim) + bias[1]
        flag = 1
        data.append([vec,flag])

    random.shuffle(data)

    x = np.array([item[0][0] for item in data])
    y = np.array([item[1] for item in data])

    return [x,y]

def normal_timeseries(length = 50, random_length = True, n_neg_sample = 100, n_pos_sample = 100, bias = [-2, +2], seed = 0):
    data = []

    random.seed(seed)
    np.random.seed(seed)

    for num in xrange(0, n_neg_sample):
        if random_length:
            s_len = np.random.randint(length) + length
        else:
            s_len = length

        sample = np.zeros(s_len)
        for i in range(0, s_len):
            sample[i] = np.random.randn(1) + bias[0]
        flag = 0
        data.append([sample, flag])

    for num in xrange(0, n_pos_sample):
        if random_length:
            s_len = np.random.randint(length) + length
        else:
            s_len = length

        sample = np.zeros(s_len)
        for i in range(0, s_len):
            sample[i] = np.random.randn(1) + bias[1]
        flag = 1
        data.append([sample, flag])

    random.shuffle(data)

    x = np.array([item[0] for item in data])
    y = np.array([item[1] for item in data])

    return [x,y]

def uniform_dist(n_dim = 2, n_sample = 100, minimum = 0.0, maximum = 1.0, seed = 1):

    np.random.seed(seed)
    return np.random.uniform(minimum, maximum, (n_sample, n_dim))
<|endoftext|>


<commit_before>#!/usr/bin/env python
from dateutil.parser import parse
import numpy as np
from matplotlib.pyplot import show
#
import pyiri90, pyiri90.plots

f0 = 3.5e6 # radar frequency [Hz]
B0 = 60e-6 # Geomagnetic field strength [Tesla] #TODO verify at altitude
theta = 0.01 # radians off parallel from B

fs = np.arange(1e6,8e6,0.05e6) # [Hz] radar frequency sweep

if __name__ == '__main__':
    from argparse import ArgumentParser
    p = ArgumentParser()
    p.add_argument('--alt',help='START STOP STEP altitude [km]',type=float,nargs=3,default=(85,500,1.))
    p.add_argument('-t','--time',help='datetime of simulation',default='2012-07-21T12:35')
    p.add_argument('-c','--latlon',help='geodetic coordinates of simulation',default=(65,-147.5),type=float)
    p.add_argument('--f107',type=float,default=200.)
    p.add_argument('--f107a',type=float,default=200.)
    p.add_argument('--ap',type=int,default=4)
    p = p.parse_args()

    altkm = np.arange(p.alt[0], p.alt[1], p.alt[2])
    dtime = parse(p.time)
#%%
    iono =  pyiri90.runiri(dtime,altkm,p.latlon,p.f107, p.f107a, ap=p.ap)

    pyiri90.plots.plotiono(iono,dtime,p.latlon,p.f107,p.f107a,p.ap)
    show()
<commit_msg>cleanup<commit_after>#!/usr/bin/env python
from dateutil.parser import parse
import numpy as np
from matplotlib.pyplot import show
#
import pyiri90, pyiri90.plots

if __name__ == '__main__':
    from argparse import ArgumentParser
    p = ArgumentParser()
    p.add_argument('--alt',help='START STOP STEP altitude [km]',type=float, nargs=3, default=(85,500,1.))
    p.add_argument('-t','--time',help='date and time of simulation',default='2012-07-21T12:35')
    p.add_argument('-c','--latlon',help='geodetic coordinates of simulation',default=(65,-147.5),type=float)
    p.add_argument('--f107',type=float,default=200.)
    p.add_argument('--f107a',type=float,default=200.)
    p.add_argument('--ap',type=int,default=4)
    p = p.parse_args()

    altkm = np.arange(p.alt[0], p.alt[1], p.alt[2])
    dtime = parse(p.time)
#%%
    iono =  pyiri90.runiri(dtime,altkm,p.latlon,p.f107, p.f107a, ap=p.ap)

    pyiri90.plots.plotiono(iono,dtime,p.latlon,p.f107,p.f107a,p.ap)
    show()
<|endoftext|>