<commit_before>import time
from scipy.spatial.distance import euclidean
from fastmc.proto import Slot
import numpy as np
import math
from math import floor
import threading
import itertools
import logging
from fastmc.proto import Position
import types

log = logging.getLogger(__name__)
log.addHandler(logging.NullHandler())

class Player(object):

    def __init__(self, proto, send_queue, recv_condition, world):
        self.proto = proto
        self._send_queue = send_queue
        self._recv_condition = recv_condition
        self.x = None
        self.y = None
        self.z = None
        self.yaw = None
        self.pitch = None
        self.health = None
        self.food = None
        self.food_saturation = 0
        self._xp_bar = -1
        self._xp_total = -1
        self.xp_level = -1
        self._available_enchantments = {}
        self._open_window_id = 0
        self._held_slot_num = 0
        self._held_slot_cycle = itertools.cycle(range(10))
        self._cursor_slot = Slot(-1, None, None, None)
        self.windows = {}
        self.world = world
        self.on_ground = True
        self._is_moving = threading.Event()
        self._position_update_lock = threading.Lock()
        self._move_lock = threading.RLock()
        self._inventory_lock = threading.RLock()
        self.move_corrected_by_server = threading.Event()

    def _wait_for(self, what, timeout=10):
        if what():
            return True
        start = time.time()
        with self._recv_condition:
            while not what() and time.time() - start < timeout:
                self._recv_condition.wait(timeout=0.1)
        return what()

    def _send(self, packet_id, **kwargs):
        self._send_queue.put((packet_id, kwargs))

    @property
    def position(self):
        return self.get_position(floor=True)

    def get_position(self, dx=0, dy=0, dz=0, floor=False):
        if self.x is None:
            return (None, None, None)
        position = np.add((self.x, self.y, self.z), (dx, dy, dz))
        if floor:
            return tuple([int(i) for i in np.floor(position)])
        else:
            return tuple(position)

    @property
    def held_item(self):
        inventory = self.inventory
        if inventory is None:
            return None
        held = inventory.held
        return held[self._held_slot_num]

    @property
    def inventory(self):
        return self.windows.get(0)

    @property
    def open_window(self):
        return self.windows.get(self._open_window_id)

    def navigate_to(self, x, y, z, speed=10, space=0, timeout=10):
        x0, y0, z0 = self.get_position(floor=True)
        x, y, z = floor(x), floor(y), floor(z)
        if euclidean((x0, y0, z0), (x, y, z)) <= space:
            return True
        log.debug('navigating from %s to %s.', str((x0, y0, z0)),
                  str((x, y, z)))
        with self._move_lock:
            path = self.world.find_path(x0, y0, z0, x, y, z, space=space,
                                        timeout=timeout)
            log.debug('path: %s', str(path))
            if not path:
                return False
            return self.follow_path(path)

    def dig_to(self, x, y, z, speed=10, space=0, timeout=10):
        x0, y0, z0 = self.get_position(floor=True)
        x, y, z = floor(x), floor(y), floor(z)
        log.debug('digging from %s to %s.', str((x0, y0, z0)), str((x, y, z)))
        if euclidean((x0, y0, z0), (x, y, z)) <= space:
            return True
        with self._move_lock:
            path = self.world.find_path(x0, y0, z0, x, y, z, space=space,
                                        timeout=timeout, digging=True)
            log.debug('path: %s', str(path))
            if not path:
                return False
            return self.follow_path(path, digging=True)

    def follow_path(self, path, speed=10, digging=False):
        log.debug('following path: %s', str(path))
        with self._move_lock:
            x0, y0, z0 = self.get_position(floor=True)
            for num, (x, y, z) in enumerate(path):
                if digging:
                    if num > 0:
                        x0, y0, z0 = path[num - 1]
                    break_set = self.world.get_blocks_to_break(
                        x0, y0, z0, x, y, z)
                    if not self.break_all_blocks(break_set):
                        log.info('could not break all the blocks: %s',
                                 str(break_set))
                        return False
                if num > 0:
                    x0, y0, z0 = path[num - 1]
                if not self.world.is_moveable(x0, y0, z0, x, y, z):
                    log.info('position is not moveable: %s',
                             str((x0, y0, z0, x, y, z)))
                    return False
                if not self.move_to(x, y, z, speed=speed, center=True):
                    log.info('could not move to: %s', str((x, y, z)))
                    return False
            return True

    def move_to(self, x, y, z, speed=10, center=False):
        def abs_min(n, delta):
            if n < 0:
                return max(n, -delta)
            else:
                return min(n, delta)

        if center:
            x = floor(x) + 0.5
            z = floor(z) + 0.5

        dt = 0.1
        delta = speed * dt
        self._is_moving.set()
        self.move_corrected_by_server.clear()
        with self._move_lock:
            while euclidean((x, y, z), self.get_position()) > 0.1:
                dx = x - self.x
                dy = y - self.y
                dz = z - self.z
                target = (abs_min(
                    dx, delta), abs_min(dy, delta), abs_min(dz, delta))
                self.move(*target)
                time.sleep(dt)
                if self.move_corrected_by_server.is_set():
                    self.move_corrected_by_server.clear()
                    self._is_moving.clear()
                    log.error("can't move from: %s to: %s", str(self.position),
                              str(target))
                    return False
        self._is_moving.clear()
        return True

    def move(self, dx=0, dy=0, dz=0):

        def calc_yaw(x0, z0, x, z):
            l = x - x0
            w = z - z0
            c = math.sqrt(l * l + w * w)
            if c == 0:
                return 0
            alpha1 = -math.asin(l / c) / math.pi * 180
            alpha2 = math.acos(w / c) / math.pi * 180
            if alpha2 > 90:
                return 180 - alpha1
            else:
                return alpha1

        with self._position_update_lock:
            self.yaw = calc_yaw(self.x, self.z, self.x + dx, self.z + dz)
            self.x += dx
            self.y += dy
            self.z += dz
            log.debug('moved to: %s', str(self.position))

    def teleport(self, x, y, z, yaw, pitch):
        with self._position_update_lock:
            self.x = x
            self.y = y
            self.z = z
            self.yaw = yaw
            self.pitch = pitch

    def iter_entities_in_range(self, types=None, reach=4):
        for entity in self.world.iter_entities(types=types):
            if euclidean((self.x, self.y, self.z),
                         (entity.x, entity.y, entity.z)) <= reach:
                yield entity

    def iter_objects_in_range(self, types=None, items=None, reach=10):
        for entity in self.world.iter_objects(types=types, items=items):
            if euclidean((self.x, self.y, self.z),
                         (entity.x, entity.y, entity.z)) <= reach:
                yield entity

    def get_closest_entity(self, types=None, limit=None):
        closest_entity, dist = None, None
        for entity in self.world.iter_entities(types=types):
            cur_dist = euclidean(self.position, (entity.x, entity.y, entity.z))
            if closest_entity is None or (cur_dist < limit and cur_dist < dist):
                closest_entity = entity
                dist = cur_dist
        return (closest_entity, dist)

    def equip_any_item_from_list(self, item_types):
        for _type in item_types:
            if self.equip_item(_type):
                return True
        return False

    def equip_item(self, item):
        with self._inventory_lock:
            if item == self.held_item:  # item already in hand
                return True
            elif item in self.inventory.held:  # item in held_items
                self.change_held_item(self.inventory.held.index(item,
                                                                relative=True))
            elif item in self.inventory:
                held_slot = self._held_slot_cycle.next()
                inventory_held_slot = len(self.inventory.slots) - 1 - held_slot
                if not self.inventory.swap_slots(inventory_held_slot,
                                                 self.inventory.index(item)):
                    return False
                held_index = self.inventory.held.index(item, relative=True)
                if held_index is None:
                    return False
                self.change_held_item(held_index)
            else:
                return False
            return self._wait_for(lambda: item == self.held_item)

    def change_held_item(self, slot_num):
        self._send(self.proto.PlayServerboundHeldItemChange.id,
                   slot=slot_num)

    def click_inventory_block(self, x, y, z):
        if self._open_window_id != 0:
            return False
        self._send(self.proto.PlayServerboundBlockPlacement.id,
                   location=Position(x, y, z),
                   direction=0,
                   held_item=None,
                   cursor_x=0,
                   cursor_y=0,
                   cursor_z=0)
        return (
            self._wait_for(lambda: self._open_window_id != 0) and
            self._wait_for(lambda: len(self.open_window.slots) > 0)
        )

    def close_window(self):
        self._send(self.proto.PlayServerboundCloseWindow.id,
                   window_id=self._open_window_id)
        for _id in self.windows.keys():
            if _id != 0:
                del self.windows[_id]
        self._open_window_id = 0

    def break_block(self, x, y, z):

        def is_changed():
            name = self.world.get_name(x, y, z)
            return name != block_name or name == 'Air'

        x, y, z = int(x), int(y), int(z)  # TODO figure out why this is needed
        log.debug('breaking block: (%d, %d, %d)', x, y, z)
        block_name = self.world.get_name(x, y, z)
        if block_name == 'Air':
            return True
        if (self.world.is_falling_block(x, y, z) or
                block_name in ('Dirt', 'Grass Block')):
            if not self.equip_any_item_from_list([
                'Diamond Shovel',
                'Golden Shovel',
                'Iron Shovel',
                'Stone Shovel',
                'Wooden Shovel',
            ]):
                log.info('No shovels to use')
        elif block_name in types.WOOD:
            if not self.equip_any_item_from_list([
                'Diamond Axe',
                'Golden Axe',
                'Iron Axe',
                'Stone Axe',
                'Wooden Axe',
            ]):
                log.info('No axes to use')
        else:
            if not self.equip_any_item_from_list([
                'Diamond Pickaxe',
                'Golden Pickaxe',
                'Iron Pickaxe',
                'Stone Pickaxe',
                'Wooden Pickaxe',
            ]):
                log.info('No pickaxes to use')
        self._send(self.proto.PlayServerboundPlayerDigging.id,
                   status=0,
                   location=Position(x, y, z),
                   face=1)
        self._send(self.proto.PlayServerboundPlayerDigging.id,
                   status=2,
                   location=Position(x, y, z),
                   face=1)
        result = self._wait_for(is_changed, timeout=5)
        if not result:
            log.info('could not break block: %s', str((x, y, z)))
        return result

    def break_all_blocks(self, blocks):
        tries = 0
        while blocks and tries < 5:
            tries += 1
            blocks = [b for b in blocks if self.world.is_solid_block(*b)]
            for block in blocks:
                self.break_block(*block)
        return not bool(blocks)
<commit_msg>added long distance navigation.<commit_after>import time
from scipy.spatial.distance import euclidean
from fastmc.proto import Slot
import numpy as np
import math
from math import floor
import threading
import itertools
import logging
from fastmc.proto import Position
import types

log = logging.getLogger(__name__)
log.addHandler(logging.NullHandler())

class Player(object):

    def __init__(self, proto, send_queue, recv_condition, world):
        self.proto = proto
        self._send_queue = send_queue
        self._recv_condition = recv_condition
        self.x = None
        self.y = None
        self.z = None
        self.yaw = None
        self.pitch = None
        self.health = None
        self.food = None
        self.food_saturation = 0
        self._xp_bar = -1
        self._xp_total = -1
        self.xp_level = -1
        self._available_enchantments = {}
        self._open_window_id = 0
        self._held_slot_num = 0
        self._held_slot_cycle = itertools.cycle(range(10))
        self._cursor_slot = Slot(-1, None, None, None)
        self.windows = {}
        self.world = world
        self.on_ground = True
        self._is_moving = threading.Event()
        self._position_update_lock = threading.Lock()
        self._move_lock = threading.RLock()
        self._inventory_lock = threading.RLock()
        self.move_corrected_by_server = threading.Event()

    def _wait_for(self, what, timeout=10):
        if what():
            return True
        start = time.time()
        with self._recv_condition:
            while not what() and time.time() - start < timeout:
                self._recv_condition.wait(timeout=0.1)
        return what()

    def _send(self, packet_id, **kwargs):
        self._send_queue.put((packet_id, kwargs))

    @property
    def position(self):
        return self.get_position(floor=True)

    def get_position(self, dx=0, dy=0, dz=0, floor=False):
        if self.x is None:
            return (None, None, None)
        position = np.add((self.x, self.y, self.z), (dx, dy, dz))
        if floor:
            return tuple([int(i) for i in np.floor(position)])
        else:
            return tuple(position)

    @property
    def held_item(self):
        inventory = self.inventory
        if inventory is None:
            return None
        held = inventory.held
        return held[self._held_slot_num]

    @property
    def inventory(self):
        return self.windows.get(0)

    @property
    def open_window(self):
        return self.windows.get(self._open_window_id)

    def navigate_to(self, x, y, z, speed=10, space=0, timeout=10):
        x0, y0, z0 = self.position
        distance = euclidean(self.position, (x, y, z))
        if distance <= space:
            return True
        log.debug('navigating from %s to %s.', str((x0, y0, z0)),
                  str((x, y, z)))
        with self._move_lock:
            while distance > space:
                space_b = min(space, 100)
                path = self.world.find_path(x0, y0, z0, x, y, z, space=space_b,
                                            timeout=timeout)
                log.debug('path: %s', str(path))
                if not path:
                    return False
                if not self.follow_path(path):
                    return False
                distance = euclidean(self.position, (x, y, z))
        return True

    def dig_to(self, x, y, z, speed=10, space=0, timeout=10):
        x0, y0, z0 = self.get_position(floor=True)
        x, y, z = floor(x), floor(y), floor(z)
        log.debug('digging from %s to %s.', str((x0, y0, z0)), str((x, y, z)))
        distance = euclidean((x0, y0, z0), (x, y, z))
        if distance <= space:
            return True
        with self._move_lock:
            while distance > space:
                space_b = min(space, 100)
                path = self.world.find_path(x0, y0, z0, x, y, z, space=space_b,
                                            timeout=timeout, digging=True)
                log.debug('path: %s', str(path))
                if not path:
                    return False
                if not self.follow_path(path, digging=True):
                    return False
                distance = euclidean(self.position, (x, y, z))
        return True

    def follow_path(self, path, speed=10, digging=False):
        log.debug('following path: %s', str(path))
        with self._move_lock:
            x0, y0, z0 = self.get_position(floor=True)
            for num, (x, y, z) in enumerate(path):
                if digging:
                    if num > 0:
                        x0, y0, z0 = path[num - 1]
                    break_set = self.world.get_blocks_to_break(
                        x0, y0, z0, x, y, z)
                    if not self.break_all_blocks(break_set):
                        log.info('could not break all the blocks: %s',
                                 str(break_set))
                        return False
                if num > 0:
                    x0, y0, z0 = path[num - 1]
                if not self.world.is_moveable(x0, y0, z0, x, y, z):
                    log.info('position is not moveable: %s',
                             str((x0, y0, z0, x, y, z)))
                    return False
                if not self.move_to(x, y, z, speed=speed, center=True):
                    log.info('could not move to: %s', str((x, y, z)))
                    return False
            return True

    def move_to(self, x, y, z, speed=10, center=False):
        def abs_min(n, delta):
            if n < 0:
                return max(n, -delta)
            else:
                return min(n, delta)

        if center:
            x = floor(x) + 0.5
            z = floor(z) + 0.5

        dt = 0.1
        delta = speed * dt
        self._is_moving.set()
        self.move_corrected_by_server.clear()
        with self._move_lock:
            while euclidean((x, y, z), self.get_position()) > 0.1:
                dx = x - self.x
                dy = y - self.y
                dz = z - self.z
                target = (abs_min(
                    dx, delta), abs_min(dy, delta), abs_min(dz, delta))
                self.move(*target)
                time.sleep(dt)
                if self.move_corrected_by_server.is_set():
                    self.move_corrected_by_server.clear()
                    self._is_moving.clear()
                    log.error("can't move from: %s to: %s", str(self.position),
                              str(target))
                    return False
        self._is_moving.clear()
        return True

    def move(self, dx=0, dy=0, dz=0):

        def calc_yaw(x0, z0, x, z):
            l = x - x0
            w = z - z0
            c = math.sqrt(l * l + w * w)
            if c == 0:
                return 0
            alpha1 = -math.asin(l / c) / math.pi * 180
            alpha2 = math.acos(w / c) / math.pi * 180
            if alpha2 > 90:
                return 180 - alpha1
            else:
                return alpha1

        with self._position_update_lock:
            self.yaw = calc_yaw(self.x, self.z, self.x + dx, self.z + dz)
            self.x += dx
            self.y += dy
            self.z += dz
            log.debug('moved to: %s', str(self.position))

    def teleport(self, x, y, z, yaw, pitch):
        with self._position_update_lock:
            self.x = x
            self.y = y
            self.z = z
            self.yaw = yaw
            self.pitch = pitch

    def iter_entities_in_range(self, types=None, reach=4):
        for entity in self.world.iter_entities(types=types):
            if euclidean((self.x, self.y, self.z),
                         (entity.x, entity.y, entity.z)) <= reach:
                yield entity

    def iter_objects_in_range(self, types=None, items=None, reach=10):
        for entity in self.world.iter_objects(types=types, items=items):
            if euclidean((self.x, self.y, self.z),
                         (entity.x, entity.y, entity.z)) <= reach:
                yield entity

    def get_closest_entity(self, types=None, limit=None):
        closest_entity, dist = None, None
        for entity in self.world.iter_entities(types=types):
            cur_dist = euclidean(self.position, (entity.x, entity.y, entity.z))
            if closest_entity is None or (cur_dist < limit and cur_dist < dist):
                closest_entity = entity
                dist = cur_dist
        return (closest_entity, dist)

    def equip_any_item_from_list(self, item_types):
        for _type in item_types:
            if self.equip_item(_type):
                return True
        return False

    def equip_item(self, item):
        with self._inventory_lock:
            if item == self.held_item:  # item already in hand
                return True
            elif item in self.inventory.held:  # item in held_items
                self.change_held_item(self.inventory.held.index(item,
                                                                relative=True))
            elif item in self.inventory:
                held_slot = self._held_slot_cycle.next()
                inventory_held_slot = len(self.inventory.slots) - 1 - held_slot
                if not self.inventory.swap_slots(inventory_held_slot,
                                                 self.inventory.index(item)):
                    return False
                held_index = self.inventory.held.index(item, relative=True)
                if held_index is None:
                    return False
                self.change_held_item(held_index)
            else:
                return False
            return self._wait_for(lambda: item == self.held_item)

    def change_held_item(self, slot_num):
        self._send(self.proto.PlayServerboundHeldItemChange.id,
                   slot=slot_num)

    def click_inventory_block(self, x, y, z):
        if self._open_window_id != 0:
            return False
        self._send(self.proto.PlayServerboundBlockPlacement.id,
                   location=Position(x, y, z),
                   direction=0,
                   held_item=None,
                   cursor_x=0,
                   cursor_y=0,
                   cursor_z=0)
        return (
            self._wait_for(lambda: self._open_window_id != 0) and
            self._wait_for(lambda: len(self.open_window.slots) > 0)
        )

    def close_window(self):
        self._send(self.proto.PlayServerboundCloseWindow.id,
                   window_id=self._open_window_id)
        for _id in self.windows.keys():
            if _id != 0:
                del self.windows[_id]
        self._open_window_id = 0

    def break_block(self, x, y, z):

        def is_changed():
            name = self.world.get_name(x, y, z)
            return name != block_name or name == 'Air'

        x, y, z = int(x), int(y), int(z)  # TODO figure out why this is needed
        log.debug('breaking block: (%d, %d, %d)', x, y, z)
        block_name = self.world.get_name(x, y, z)
        if block_name == 'Air':
            return True
        if (self.world.is_falling_block(x, y, z) or
                block_name in ('Dirt', 'Grass Block')):
            if not self.equip_any_item_from_list([
                'Diamond Shovel',
                'Golden Shovel',
                'Iron Shovel',
                'Stone Shovel',
                'Wooden Shovel',
            ]):
                log.info('No shovels to use')
        elif block_name in types.WOOD:
            if not self.equip_any_item_from_list([
                'Diamond Axe',
                'Golden Axe',
                'Iron Axe',
                'Stone Axe',
                'Wooden Axe',
            ]):
                log.info('No axes to use')
        else:
            if not self.equip_any_item_from_list([
                'Diamond Pickaxe',
                'Golden Pickaxe',
                'Iron Pickaxe',
                'Stone Pickaxe',
                'Wooden Pickaxe',
            ]):
                log.info('No pickaxes to use')
        self._send(self.proto.PlayServerboundPlayerDigging.id,
                   status=0,
                   location=Position(x, y, z),
                   face=1)
        self._send(self.proto.PlayServerboundPlayerDigging.id,
                   status=2,
                   location=Position(x, y, z),
                   face=1)
        result = self._wait_for(is_changed, timeout=5)
        if not result:
            log.info('could not break block: %s', str((x, y, z)))
        return result

    def break_all_blocks(self, blocks):
        tries = 0
        while blocks and tries < 5:
            tries += 1
            blocks = [b for b in blocks if self.world.is_solid_block(*b)]
            for block in blocks:
                self.break_block(*block)
        return not bool(blocks)
<|endoftext|>


<commit_before>#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""Machine learning model which uses Wikipedia data to predicts ILI levels.

Usage:
  model.py <year_start> <year_end> <dataset_path> <incidence_path> <keywords_file> <country_name> [--p] [--f] [--v] [--d=<directory>] [--no-future] [--no-images]

  <year_start>      The first influenza season we want to predict.
  <year_end>        The last influenza season we want to predict.
  <dataset_path>    The path to dataset files (the ones which contain Wikipedia pageviews).
  <incidence_path>  The path to the files which specify the ILI levels for each week.
  <keywords_file>   The path to the file which contains the Wikipedia pages used.
  <country_name>    The name of the country we are analyzing
  -p, --poisson     Use the Poisson model + LASSO instead of the linear one.
  -f, --file        Write informations to file
  -v, --verbose     Output more informations
  -d, --directory   Select a directory in which save your files
  -n, --no-future   Use a different method to train the model (avoid using seasonal influenza which are
                    later than the one we want to predict)
  --no-images       Do not generate any graphs. Print only to screen.
  -h, --help        Print this help message
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import csv
from sklearn.linear_model import LassoCV
from sklearn.metrics import mean_squared_error
from tabulate import tabulate
from docopt import docopt

from models_utils import generate, generate_labels, generate_labels_one_year, generate_one_year, standardize_data, generate_keywords, stz, get_important_pages, correlation_matrix

from cvglmnetCoef import cvglmnetCoef
from cvglmnetPredict import cvglmnetPredict
from cvglmnet import cvglmnet

#########################
#### INITIALIZATION #####
#########################

# Parse the command line
arguments = docopt(__doc__)

# Directory where to save files
save_directory="./"
if arguments["--d"]:
    save_directory = arguments["--d"]

# Year selected
year_start = int(arguments["<year_start>"])
year_end = int(arguments["<year_end>"])+1

# Model type
model_type = "Linear Model"
using_poisson = False
if arguments["--p"]:
    using_poisson = True
    model_type="Poisson Model"

# Feature path and labels
path_features = arguments["<dataset_path>"]
path_labels = arguments["<incidence_path>"]
keywords = arguments["<keywords_file>"]

# Selected columns that will be extracted from the dataframes
selected_columns = generate_keywords(keywords)

all_features_values = pd.DataFrame()
all_predicted_values = []
all_true_labels = pd.DataFrame()
total_weeks = []
all_weighted_feature=dict()

# The pair start-end seasons
year_sel = (year_start, year_end)

# Useful information which can be written to file
peaks = []
predicted_peaks = []
peaks_value = []
predicted_peaks_value = []
mse_seasons = []

# If --f then we write some information to file
if arguments["--f"]:
    csvfile = open(save_directory+str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+"_information_"+arguments["<country_name>"]+".csv", 'w', newline='')
    spamwriter = csv.writer(csvfile, delimiter=',', quotechar='|', quoting=csv.QUOTE_MINIMAL)
    spamwriter.writerow(['season', 'country', 'predicted_influenza_peak', 'real_influenza_peak', 'predicted_influenza_peak_value', 'real_influenza_peak_value', 'mse', 'pcc'])

######################
##### ALGORITHM ######
######################

# For each season, train a model with the complete dataset (without the season)
# and predict the ILI incidence.
for year_selected in range(year_sel[0], year_sel[1]):

    print("------------")
    print("[*] ", year_selected )

    # Generate the list of excluded years
    excluded=[]
    if arguments["--no-future"]:
        for i in range(2007, 2018):
            if i>year_selected:
                excluded.append(i)

    # Data generation from data files
    dataset = generate(year_selected, excluded, path_features)[selected_columns]
    labels = generate_labels(year_selected, excluded, path_labels)
    data = generate_one_year(year_selected, path_features)[selected_columns]
    labels_test = generate_labels_one_year(year_selected, path_labels)
    weeks = labels_test["week"]

    # Create a copy of the dataset where we fills NaN values
    # with 0.
    dataset_zero = dataset.fillna(0)
    data_zero = data.fillna(0)

    # Standardize data
    train, test = standardize_data(dataset_zero, data_zero)

    # Create a Lasso Cross-Validation instance which will be
    # trained on the dataset in which NaN values are replaced
    # with the column mean.
    if not using_poisson:
        model = LassoCV(max_iter=100000, n_jobs=-1, n_alphas=1000, random_state=1)
        model.fit(train, labels["incidence"].fillna(0))
        result = model.predict(test)
    else:
        model = cvglmnet(x=train.values.copy(), y=labels["incidence"].fillna(0).copy().values, family='poisson', alpha=1.0,
                       ptype="mse", parallel=True, nfolds=10)
        result = cvglmnetPredict(model, test.values, ptype='response', s="lambda_min")
        coeff = cvglmnetCoef(model, s="lambda_min")

    # Get the feature coefficients
    if not using_poisson:
        coeff = model.coef_
    else:
        coeff = cvglmnetCoef(model, s="lambda_min")

    # Add the pair (value, coeff) to a dictionary
    for i in list(zip(coeff, selected_columns)):
        if (i[0] != 0):
            if all_weighted_feature.get(i[1], None) is None:
                all_weighted_feature[i[1]] = [i[0]]
            else:
                all_weighted_feature[i[1]].append(i[0])

    # Add the value to global variables
    all_features_values= all_features_values.append(data.fillna(0))
    all_predicted_values.extend(result)
    all_true_labels = all_true_labels.append(labels_test)
    total_weeks.extend(weeks)

    mse_seasons.append(mean_squared_error(labels_test['incidence'].fillna(0), result))

    # If we are verbose print more informations
    index_peak = labels_test['incidence'].fillna(0).idxmax(axis=1)
    index_peak_m = np.argmax(result)
    peaks.append(labels_test.iloc[index_peak]["week"])
    predicted_peaks.append(labels_test.iloc[index_peak_m]["week"])
    peaks_value.append(labels_test.iloc[index_peak]["incidence"])
    predicted_peaks_value.append(result[index_peak_m])
    if arguments["--v"]:
        print("[*] Influenza Season Peak (Week):", labels_test.iloc[index_peak]["week"])
        print("[*] Influenza Seasons Predicted Peak (Week): ", labels_test.iloc[index_peak_m]["week"])
        print("[*] Influenza Season Peak Value: ", labels_test.iloc[index_peak]["incidence"])
        print("[*] Influenza Season Predicted Peak Value: ", result[index_peak_m])
        print("------------")

    # For each influenza seasons we print data to file
    if arguments["--f"]:
        spamwriter.writerow([str(year_sel[0]-1)+"-"+str(year_sel[1]-1),
                             arguments["<country_name>"],
                             labels_test.iloc[index_peak_m]["week"],
                             labels_test.iloc[index_peak]["week"],
                             result[index_peak_m],
                             labels_test.iloc[index_peak]["incidence"],
                             mean_squared_error(labels_test['incidence'].fillna(0), result),
                             np.corrcoef(result, labels_test['incidence'].fillna(0), rowvar=False)[0][1]
                             ])

# Get important pages to generate the plot we need
important_pages = get_important_pages(all_weighted_feature, 5)

# Print MSE of the two models
mse = mean_squared_error(all_true_labels["incidence"].fillna(0), all_predicted_values)
print("------------")
print("MSE: ", mse)
print("------------")

# Print Pearson Coefficient
pcc = np.corrcoef(all_predicted_values, all_true_labels["incidence"].fillna(0), rowvar=False)[0][1]
print("------------")
print("Pearson Correlation Coeff: ", pcc)
print("------------")

# Print important pages
print(tabulate(important_pages, headers=["Page name", "Weight"]))

###########################
#### GRAPHS GENERATION ####
###########################

if not arguments["--no-images"]:
    # Plot some informations
    fig = plt.figure(3, figsize=(15, 6))
    ax = fig.add_subplot(111)

    # Set up the graph title, x and y axis titles
    plt.title("Influenza Seasons "+str(year_sel[0]-1)+" - "+str(year_sel[1]-1), fontsize=18)
    plt.ylabel("Incidence on 100000 people", fontsize=17)
    plt.xlabel("Year-Week", fontsize=17)

    # Generate the x axis labels in such way to
    # have the year-week pair only every two ticks.
    weeks_used = []
    for k, v in enumerate(total_weeks):
        if k%2==0:
            weeks_used.append(v)
        else:
            weeks_used.append(" ")

    # Set up the axes ticks
    plt.xticks(range(0, len(weeks_used)), weeks_used, rotation="vertical", fontsize=15)
    plt.yticks(fontsize=15)

    # Plot the model result and the incidence
    plt.plot(range(0, len(all_predicted_values)), all_predicted_values, 'r-', label=model_type, linewidth=3)
    plt.plot(range(0, len(all_true_labels["incidence"])), all_true_labels["incidence"].fillna(0), 'k-', label="Incidence", linewidth=3)

    # Add dotted line into the graph to delimit the influenza seasons.
    for i in range(1, (year_sel[1]-year_sel[0])):
        plt.axvline(26*i, linestyle="dashed", color="k", linewidth=3)

    # Update some graphical options
    plt.grid()
    plt.legend(fontsize=16, loc="upper right")
    plt.tight_layout()

    # Save the graph on file
    plt.savefig(save_directory+"season_"+str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+".png", dpi=150)

    # Feature plot
    figf = plt.figure(4, figsize=(15, 6))
    axf = figf.add_subplot(111)

    # Generate the graph showing the pageview variation of the top 5 Wikipedia's pages
    # choosen by the model.
    plt.title("Pageview variation "+str(year_sel[0]-1)+" - "+str(year_sel[1]-1), fontsize=18)
    plt.ylabel("Pageview value", fontsize=17)
    plt.xlabel("Year-Week", fontsize=17)

    # Plot the axes labels
    plt.xticks(range(0, len(weeks_used)), weeks_used, rotation="vertical", fontsize=15)
    plt.yticks(fontsize=15)

    # Plot all the pageview data and the ILI incidence
    std_all_features_values = stz(all_features_values)
    for key, value in important_pages:
     plt.plot(range(0, len(std_all_features_values[key])), std_all_features_values[key], '-', label=key, linewidth=3)
    plt.plot(range(0, len(all_true_labels["incidence"])), stz(all_true_labels["incidence"]), 'k-', label="Incidence", linewidth=3)

    # Add dotted line into the graph to delimit the influenza seasons.
    for i in range(1, (year_sel[1]-year_sel[0])):
        plt.axvline(26*i, linestyle="dashed", color="k", linewidth=3)

    # Update some graphical options
    plt.grid()
    plt.legend(fontsize=16, loc="best")
    plt.tight_layout()

    # Save again the graph on file
    plt.savefig(save_directory+"season_"+str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+"_features.png", dpi=150)

    # Correlation Matrix
    new_pages=[]
    for p in important_pages:
        new_pages.append(p[0])
    first = pd.DataFrame(all_features_values[new_pages])
    second = pd.DataFrame(all_true_labels["incidence"]).set_index(first.index.values)
    first['incidence'] = second
    first = first.loc[:, (first != 0).any(axis=0)]
    new_pages = list(first.columns.values)
    correlation_matrix(first, "Correlation Matrix "+str(year_sel[0]-1)+"-"+str(year_sel[1]-1), new_pages, save_directory+"cmatrix_"+str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+".png")

###################
#  SAVE TO FILE   #
###################

# If --f then we write some information to file
#if arguments["--f"]:
#    with open(save_directory+str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+"_information.csv", 'w', newline='') as csvfile:
#        spamwriter = csv.writer(csvfile, delimiter=',',
#                                quotechar='|', quoting=csv.QUOTE_MINIMAL)
#        spamwriter.writerow(['real_influenza_peak', peaks])
#        spamwriter.writerow(['predicted_influenza_peak', predicted_peaks])
#        spamwriter.writerow(['real_influenza_peak_value', peaks_value])
#        spamwriter.writerow(['predicted_influenza_peak_value', predicted_peaks_value])
#        spamwriter.writerow(['mse_all_seasons', mse_seasons])
#        spamwriter.writerow(['mse', mean_squared_error(all_true_labels["incidence"].fillna(0), all_predicted_values)])
#        spamwriter.writerow(['pearson', pcc])
#        for p in important_pages:
#            spamwriter.writerow([p[0], float(p[1])])

<commit_msg>[model] Print feature information to file in a different way.<commit_after>#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""Machine learning model which uses Wikipedia data to predicts ILI levels.

Usage:
  model.py <year_start> <year_end> <dataset_path> <incidence_path> <keywords_file> <country_name> [--p] [--f] [--v] [--d=<directory>] [--no-future] [--no-images]

  <year_start>      The first influenza season we want to predict.
  <year_end>        The last influenza season we want to predict.
  <dataset_path>    The path to dataset files (the ones which contain Wikipedia pageviews).
  <incidence_path>  The path to the files which specify the ILI levels for each week.
  <keywords_file>   The path to the file which contains the Wikipedia pages used.
  <country_name>    The name of the country we are analyzing
  -p, --poisson     Use the Poisson model + LASSO instead of the linear one.
  -f, --file        Write informations to file
  -v, --verbose     Output more informations
  -d, --directory   Select a directory in which save your files
  -n, --no-future   Use a different method to train the model (avoid using seasonal influenza which are
                    later than the one we want to predict)
  --no-images       Do not generate any graphs. Print only to screen.
  -h, --help        Print this help message
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import csv
from sklearn.linear_model import LassoCV
from sklearn.metrics import mean_squared_error
from tabulate import tabulate
from docopt import docopt

from models_utils import generate, generate_labels, generate_labels_one_year, generate_one_year, standardize_data, generate_keywords, stz, get_important_pages, correlation_matrix

from cvglmnetCoef import cvglmnetCoef
from cvglmnetPredict import cvglmnetPredict
from cvglmnet import cvglmnet

#########################
#### INITIALIZATION #####
#########################

# Parse the command line
arguments = docopt(__doc__)

# Directory where to save files
save_directory="./"
if arguments["--d"]:
    save_directory = arguments["--d"]

# Year selected
year_start = int(arguments["<year_start>"])
year_end = int(arguments["<year_end>"])+1

# Model type
model_type = "Linear Model"
using_poisson = False
if arguments["--p"]:
    using_poisson = True
    model_type="Poisson Model"

# Feature path and labels
path_features = arguments["<dataset_path>"]
path_labels = arguments["<incidence_path>"]
keywords = arguments["<keywords_file>"]

# Selected columns that will be extracted from the dataframes
selected_columns = generate_keywords(keywords)

all_features_values = pd.DataFrame()
all_predicted_values = []
all_true_labels = pd.DataFrame()
total_weeks = []
all_weighted_feature=dict()

# The pair start-end seasons
year_sel = (year_start, year_end)

# Useful information which can be written to file
peaks = []
predicted_peaks = []
peaks_value = []
predicted_peaks_value = []
mse_seasons = []

# If --f then we write some information to file
if arguments["--f"]:
    csvfile = open(save_directory+str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+"_information_"+arguments["<country_name>"]+".csv", 'w', newline='')
    spamwriter = csv.writer(csvfile, delimiter=',', quotechar='|', quoting=csv.QUOTE_MINIMAL)
    spamwriter.writerow(['season', 'country', 'predicted_influenza_peak', 'real_influenza_peak', 'predicted_influenza_peak_value', 'real_influenza_peak_value', 'mse', 'pcc'])

    top_features_file = open(save_directory+str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+"_features_"+arguments["<country_name>"]+".csv", 'w', newline='')
    top_features_writer = csv.writer(top_features_file, delimiter=',', quotechar='|', quoting=csv.QUOTE_MINIMAL)
    top_features_writer.writerow(['season', 'country', 'page_name', 'value'])

######################
##### ALGORITHM ######
######################

# For each season, train a model with the complete dataset (without the season)
# and predict the ILI incidence.
for year_selected in range(year_sel[0], year_sel[1]):

    print("------------")
    print("[*] ", year_selected )

    # Generate the list of excluded years
    excluded=[]
    if arguments["--no-future"]:
        for i in range(2007, 2018):
            if i>year_selected:
                excluded.append(i)

    # Data generation from data files
    dataset = generate(year_selected, excluded, path_features)[selected_columns]
    labels = generate_labels(year_selected, excluded, path_labels)
    data = generate_one_year(year_selected, path_features)[selected_columns]
    labels_test = generate_labels_one_year(year_selected, path_labels)
    weeks = labels_test["week"]

    # Create a copy of the dataset where we fills NaN values
    # with 0.
    dataset_zero = dataset.fillna(0)
    data_zero = data.fillna(0)

    # Standardize data
    train, test = standardize_data(dataset_zero, data_zero)

    # Create a Lasso Cross-Validation instance which will be
    # trained on the dataset in which NaN values are replaced
    # with the column mean.
    if not using_poisson:
        model = LassoCV(max_iter=100000, n_jobs=-1, n_alphas=1000, random_state=1)
        model.fit(train, labels["incidence"].fillna(0))
        result = model.predict(test)
    else:
        model = cvglmnet(x=train.values.copy(), y=labels["incidence"].fillna(0).copy().values, family='poisson', alpha=1.0,
                       ptype="mse", parallel=True, nfolds=10)
        result = cvglmnetPredict(model, test.values, ptype='response', s="lambda_min")
        coeff = cvglmnetCoef(model, s="lambda_min")

    # Get the feature coefficients
    if not using_poisson:
        coeff = model.coef_
    else:
        coeff = cvglmnetCoef(model, s="lambda_min")

    # Add the pair (value, coeff) to a dictionary
    current_weighted_features = dict()
    for i in list(zip(coeff, selected_columns)):
        if (i[0] != 0):
            if all_weighted_feature.get(i[1], None) is None:
                all_weighted_feature[i[1]] = [i[0]]
            else:
                all_weighted_feature[i[1]].append(i[0])
            if current_weighted_features.get(i[1], None) is None:
                current_weighted_features[i[1]] = [i[0]]
            else:
                current_weighted_features[i[1]].append(i[0])

    # Add the value to global variables
    all_features_values= all_features_values.append(data.fillna(0))
    all_predicted_values.extend(result)
    all_true_labels = all_true_labels.append(labels_test)
    total_weeks.extend(weeks)

    mse_seasons.append(mean_squared_error(labels_test['incidence'].fillna(0), result))

    # If we are verbose print more informations
    index_peak = labels_test['incidence'].fillna(0).idxmax(axis=1)
    index_peak_m = np.argmax(result)
    peaks.append(labels_test.iloc[index_peak]["week"])
    predicted_peaks.append(labels_test.iloc[index_peak_m]["week"])
    peaks_value.append(labels_test.iloc[index_peak]["incidence"])
    predicted_peaks_value.append(result[index_peak_m])
    if arguments["--v"]:
        print("[*] Influenza Season Peak (Week):", labels_test.iloc[index_peak]["week"])
        print("[*] Influenza Seasons Predicted Peak (Week): ", labels_test.iloc[index_peak_m]["week"])
        print("[*] Influenza Season Peak Value: ", labels_test.iloc[index_peak]["incidence"])
        print("[*] Influenza Season Predicted Peak Value: ", result[index_peak_m])
        print("------------")

    # For each influenza seasons we print data to file
    if arguments["--f"]:
        spamwriter.writerow([str(year_selected-1)+"-"+str(year_selected),
                             arguments["<country_name>"],
                             labels_test.iloc[index_peak_m]["week"],
                             labels_test.iloc[index_peak]["week"],
                             result[index_peak_m],
                             labels_test.iloc[index_peak]["incidence"],
                             mean_squared_error(labels_test['incidence'].fillna(0), result),
                             np.corrcoef(result, labels_test['incidence'].fillna(0), rowvar=False)[0][1]
                             ])
        for p in get_important_pages(current_weighted_features, 10):
            top_features_writer.writerow([str(year_selected-1)+"-"+str(year_selected),
                                          arguments["<country_name>"],
                                          p[0],
                                          float(p[1])])

# Get important pages to generate the plot we need
important_pages = get_important_pages(all_weighted_feature, 5)

# Print MSE of the two models
mse = mean_squared_error(all_true_labels["incidence"].fillna(0), all_predicted_values)
print("------------")
print("MSE: ", mse)
print("------------")

# Print Pearson Coefficient
pcc = np.corrcoef(all_predicted_values, all_true_labels["incidence"].fillna(0), rowvar=False)[0][1]
print("------------")
print("Pearson Correlation Coeff: ", pcc)
print("------------")

# Print important pages
print(tabulate(important_pages, headers=["Page name", "Weight"]))

###########################
#### GRAPHS GENERATION ####
###########################

if not arguments["--no-images"]:
    # Plot some informations
    fig = plt.figure(3, figsize=(15, 6))
    ax = fig.add_subplot(111)

    # Set up the graph title, x and y axis titles
    plt.title("Influenza Seasons "+str(year_sel[0]-1)+" - "+str(year_sel[1]-1), fontsize=18)
    plt.ylabel("Incidence on 100000 people", fontsize=17)
    plt.xlabel("Year-Week", fontsize=17)

    # Generate the x axis labels in such way to
    # have the year-week pair only every two ticks.
    weeks_used = []
    for k, v in enumerate(total_weeks):
        if k%2==0:
            weeks_used.append(v)
        else:
            weeks_used.append(" ")

    # Set up the axes ticks
    plt.xticks(range(0, len(weeks_used)), weeks_used, rotation="vertical", fontsize=15)
    plt.yticks(fontsize=15)

    # Plot the model result and the incidence
    plt.plot(range(0, len(all_predicted_values)), all_predicted_values, 'r-', label=model_type, linewidth=3)
    plt.plot(range(0, len(all_true_labels["incidence"])), all_true_labels["incidence"].fillna(0), 'k-', label="Incidence", linewidth=3)

    # Add dotted line into the graph to delimit the influenza seasons.
    for i in range(1, (year_sel[1]-year_sel[0])):
        plt.axvline(26*i, linestyle="dashed", color="k", linewidth=3)

    # Update some graphical options
    plt.grid()
    plt.legend(fontsize=16, loc="upper right")
    plt.tight_layout()

    # Save the graph on file
    plt.savefig(save_directory+"season_"+str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+".png", dpi=150)

    # Feature plot
    figf = plt.figure(4, figsize=(15, 6))
    axf = figf.add_subplot(111)

    # Generate the graph showing the pageview variation of the top 5 Wikipedia's pages
    # choosen by the model.
    plt.title("Pageview variation "+str(year_sel[0]-1)+" - "+str(year_sel[1]-1), fontsize=18)
    plt.ylabel("Pageview value", fontsize=17)
    plt.xlabel("Year-Week", fontsize=17)

    # Plot the axes labels
    plt.xticks(range(0, len(weeks_used)), weeks_used, rotation="vertical", fontsize=15)
    plt.yticks(fontsize=15)

    # Plot all the pageview data and the ILI incidence
    std_all_features_values = stz(all_features_values)
    for key, value in important_pages:
     plt.plot(range(0, len(std_all_features_values[key])), std_all_features_values[key], '-', label=key, linewidth=3)
    plt.plot(range(0, len(all_true_labels["incidence"])), stz(all_true_labels["incidence"]), 'k-', label="Incidence", linewidth=3)

    # Add dotted line into the graph to delimit the influenza seasons.
    for i in range(1, (year_sel[1]-year_sel[0])):
        plt.axvline(26*i, linestyle="dashed", color="k", linewidth=3)

    # Update some graphical options
    plt.grid()
    plt.legend(fontsize=16, loc="best")
    plt.tight_layout()

    # Save again the graph on file
    plt.savefig(save_directory+"season_"+str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+"_features.png", dpi=150)

    # Correlation Matrix
    new_pages=[]
    for p in important_pages:
        new_pages.append(p[0])
    first = pd.DataFrame(all_features_values[new_pages])
    second = pd.DataFrame(all_true_labels["incidence"]).set_index(first.index.values)
    first['incidence'] = second
    first = first.loc[:, (first != 0).any(axis=0)]
    new_pages = list(first.columns.values)
    correlation_matrix(first, "Correlation Matrix "+str(year_sel[0]-1)+"-"+str(year_sel[1]-1), new_pages, save_directory+"cmatrix_"+str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+".png")

###################
#  SAVE TO FILE   #
###################

# If --f then we write some information to file
#if arguments["--f"]:
#    with open(save_directory+str(year_sel[0]-1)+"-"+str(year_sel[1]-1)+"_information.csv", 'w', newline='') as csvfile:
#        spamwriter = csv.writer(csvfile, delimiter=',',
#                                quotechar='|', quoting=csv.QUOTE_MINIMAL)
#        spamwriter.writerow(['real_influenza_peak', peaks])
#        spamwriter.writerow(['predicted_influenza_peak', predicted_peaks])
#        spamwriter.writerow(['real_influenza_peak_value', peaks_value])
#        spamwriter.writerow(['predicted_influenza_peak_value', predicted_peaks_value])
#        spamwriter.writerow(['mse_all_seasons', mse_seasons])
#        spamwriter.writerow(['mse', mean_squared_error(all_true_labels["incidence"].fillna(0), all_predicted_values)])
#        spamwriter.writerow(['pearson', pcc])
#        for p in important_pages:
#            spamwriter.writerow([p[0], float(p[1])])

<|endoftext|>


<commit_before>
#!/usr/bin/env python2.7

"""%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

AGNfitter

==============================

Fitting SEDs of AGN in a MCMC Approach
G.Calistro Rivera, E.Lusso, J.Hennawi, D. Hogg

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

This is the main script.

For default use
change only the functions which state
***USER INPUT NEEDED***.

"""

#PYTHON IMPORTS
import sys,os
import numpy as np
import matplotlib.pyplot as plt
import pylab as pl
import time
import shelve
import multiprocessing as mp
import cPickle
import argparse

#AGNfitter IMPORTS
from functions import  MCMC_AGNfitter, PLOTandWRITE_AGNfitter
import functions.PARAMETERSPACE_AGNfitter as parspace
from functions.DATA_AGNfitter import DATA, DATA_all
from functions.PLOTandWRITE_AGNfitter import CHAIN, FLUXES_ARRAYS
from functions.DICTIONARIES_AGNfitter import MODELSDICT
import functions.CONSTRUCT_modelobjects as MODELFILES
from astropy import units as u

def CATALOG_settings():

    """==================================
    ***USER INPUT NEEDED***

    Set the right values to be able to read your catalog's format.
    FITS option is not available yet.
    =================================="""

    cat = dict()

    ##GENERAL
    cat['path'] ='/Users/USER/AGNfitter/'  #YOUR OWN PATH

    cat['filename'] = 'data/catalog_example.txt'
    cat['filetype'] = 'ASCII' ## catalog file type: 'ASCII' or 'FITS'.
                              ## FITS option not available yet.

    cat['name'] = 0#'ID'            ## If ASCII: Column index (int) of source IDs
                                        ## If FITS: not yet
    cat['redshift'] = 1#'z'              ## If ASCII:  Column index(int) of redshift
                                        ## If FITS: not yet

    ##FREQUENCIES/WAVELENGTHS
    ## if ASCII specify 'freq/wl_list', if FITS specify 'freq/wl_suffix'
    cat['freq/wl_list'] = np.arange(2,56,3).tolist()
                                        ## If ASCII: List of column indexes (int),
                                        ##           corresponding to freq/wl.
    #cat['freq/wl_suffix'] = '_wl'      ## If FITS: common ending to wavelength column names
    cat['freq/wl_format'] = 'wavelength' ## Gives your catalog *observed*
                                         ## 'frequency' or 'wavelength'?
    cat['freq/wl_unit'] = u.Angstrom       ## Astropy unit of freq or wavelength

    ##FLUXES
    ## if ASCII specify 'freq/wl_list', if FITS specify 'freq/wl_suffix'
    cat['flux_unit'] = u.Jy             ## Astropy unit of *flux* (astropy-units)
    cat['flux_list'] = np.arange(3,57,3).tolist()
                                        ## If ASCII: List of column indexes (int)
    #cat['flux_suffix'] = '_f'          ## If FITS: Common ending of all flux column names (str)

    cat['fluxerr_list'] = np.arange(4,58,3).tolist()
                                        ## If ASCII: List of column indexes (int)
    #cat['fluxerr_suffix'] = '_e'       ## If FITS: common ending to fluxerr column names (str)

    ##NON-DETECTIONS
    cat['ndflag_bool'] = False          ## Does you catalog has columns with flags 1(0) for
                                        ## detections (nondetections)?
    cat['ndflag_list'] = 'list'         ## If ASCII: List of column indexes (int)
                                        ## If FITS: List of column names (str)

    ##OTHERS (user input ***not*** needed)
    cat['output_folder'] =  cat['path'] +'OUTPUT/'#if no special OUTPUT folder, leave default
    cat['dict_path'] = cat['path'] + 'models/MODELSDICT_default'

    return cat

def FILTERS_settings():

    """==================================
    Set the photometric bands included in your catalog,
    in order to integrate the models over their response curves.
    =================================="""

    filters = dict()

    filters['dict_zarray'] =np.array([0.283, 1.58])  # The grid of redshifts needed to fit your catalog
    filters['Bandset'] = 'BANDSET_default' # OPTIONS:
                                           # 'BANDSET_default' (for testing)
                                           # 'BANDSET_settings' (choosing relevant filters below, as given by your catalog)
                                           # if your filter is not included, go to DICTIONARIES_AGNfitter to add.

    filters['SPIRE500']= True
    filters['SPIRE350']= True
    filters['SPIRE250']= True
    filters['PACS160']=False
    filters['PACS100']=False

    filters['MIPS160']=False
    filters['MIPS70']=False
    filters['MIPS24']=True

    filters['IRAC4']=True
    filters['IRAC3']=True
    filters['IRAC2']=True
    filters['IRAC1']=True

    filters['WISE4']=False
    filters['WISE3']=False
    filters['WISE2']=False
    filters['WISE1']=False

    filters['Ks_2mass']=True
    filters['H_2mass']=True
    filters['J_2mass']=True

    filters['H_VISTA']=False
    filters['J_VISTA']=False
    filters['K_VISTA']=False
    filters['Y_VISTA']=True

    filters['u_SDSS']=False
    filters['g_SDSS']=False
    filters['r_SDSS']=False
    filters['i_SDSS']=False
    filters['z_SDSS']=False

    filters['g_SUBARU']=False
    filters['r_SUBARU']=True
    filters['i_SUBARU']=True
    filters['z_SUBARU']=True
    filters['B_SUBARU']=True
    filters['V_SUBARU']=False

    filters['u_CHFT']=True
    filters['g_CHFT']=False
    filters['r_CHFT']=False
    filters['i_CHFT']=False
    filters['z_CHFT']=False

    filters['GALEX_2500']=True
    filters['GALEX_1500']=False

    return filters

def MCMC_settings():

    """==================================
    Set your preferences for the MCMC sampling.
    =================================="""

    mc = dict()

    mc['Nwalkers'] = 100  ## number of walkers
    mc['Nburnsets']= 2   ## number of burn-in sets
    mc['Nburn'] = 4000 ## length of each burn-in sets
    mc['Nmcmc'] = 10000  ## length of each burn-in sets
    mc['iprint'] = 1000 ## show progress in terminal in steps of this many samples

    return mc

def OUTPUT_settings():

    """==================================
    Set your preferences for the production of OUTPUT files.
    =================================="""

    out = dict()

    out['plot_format'] = 'pdf'

    #CHAIN TRACES
    out['plot_tracesburn-in'] = False
    out['plot_tracesmcmc'] = True

    #BASIC OUTPUT
    out['Nsample'] = 1000 ## out['Nsample'] * out['Nthinning'] <= out['Nmcmc']
    out['Nthinning'] = 10 ## This describes thinning of the chain to sample
    out['writepar_meanwitherrors'] = True ##Write output values for all parameters in a file.
    out['plot_posteriortriangle'] = False ##Plot triangle with all parameters' PDFs?

    #INTEGRATED LUMINOSITIES
    out['calc_intlum'] = True
    out['realizations2int'] = 100 #This process is very time consuming.
                                #Around 100-1000 is recomendend for computational reasons.
                                #If you want to plot posterior triangles of
                                #the integrated luminosities, should be > 1000.
    out['plot_posteriortrianglewithluminosities'] = False  # requires out['calc_intlum']=True

    #INTEGRATION RANGES
    out['intlum_models'] = ['sb','bbb', 'bbbdered', 'gal', 'tor','sb']  #leave 'sb' always
                                                                        #as first element
    out['intlum_freqranges_unit'] = u.micron   #Astropy unit
    out['intlum_freqranges'] = np.array([[8.,1000.],[0.1,1.],[0.1,1.],[0.1,1.],[1.,30.],[1.,30.]])
    out['intlum_names'] = ['LIR(8-1000)','Lbb(0.1-1)', 'Lbbdered(0.1-1)', 'Lga(01-1)', 'Ltor(1-30)','Lsb(1-30)']

    #SED PLOTTING
    out['realizations2plot'] = 10

    out['plotSEDrealizations'] = True

    return out

def header():
    print '              '
    print '             XXXX'
    print '___________ XXX _________________________________________________'
    print '            XX      '
    print '            X     '
    print '            X                       AGNfitter                     '
    print '         __ X __                    ---------                '
    print '     /**\   |   /**\                                          '
    print '... (*** =  o  = ***) ...........................................'
    print '     \**/__ | __\**/                                     '
    print '            X              Fitting SEDs of AGN and Galaxies  '
    print '            X             in a MCMC Approach '
    print '           xx              (Calistro Rivera et al. 2016)    '
    print '          xx               '
    print '_______ xxx______________________________________________________'
    print '     xxxx'
    print ''
    return

def MAKE_model_dictionary(cat, clobbermodel=False):
    ## 0. CONSTRUCT DICTIONARY (not needed if default is used)
    t0= time.time()

    if clobbermodel and os.path.lexists(cat['dict_path']):
        print "removing model dictionary "+cat['dict_path']
        os.system('rm -rf '+ cat['dict_path'])

    if not os.path.lexists(cat['dict_path']):

        MODELFILES.construct(cat['path'])

        mydict = MODELSDICT(cat['dict_path'], cat['path'], filters)
        mydict.build()

	print '_____________________________________________________'
	print 'For this dictionary creation %.2g min elapsed'% ((time.time() - t0)/60.)

    Modelsdict = cPickle.load(file(cat['dict_path'], 'rb'))

    return Modelsdict

def RUN_AGNfitter_onesource_independent( line, data_obj, modelsdict, clobbermodel=False):
    """
    Main function for fitting a single source in line 'line' and create it's modelsdict independently.
    """

    mc = MCMC_settings()
    out = OUTPUT_settings()

    data = DATA(data_obj,line)
    data.DICTS(filters, Modelsdict)

    print ''
    print 'Fitting sources from catalog: ', data.catalog
    print '- Sourceline: ', line
    print '- Sourcename: ', data.name

    ## 0. CONSTRUCT DICTIONARY for this redshift
    t0= time.time()

    filtersz = FILTERS_settings([data.z])

    # add a suffix for this source dictionary
    dictz = cat['dict_path'] + '_' + str(data.name)
    if clobbermodel and os.path.lexists(dictz):
        os.system('rm -rf '+dictz)
        print "removing source model dictionary "+dictz

    if not os.path.lexists(dictz):
        mydict = MODELSDICT(dictz, cat['path'], filtersz)
        mydict.build()
	print '_____________________________________________________'
	print 'For this dictionary creation %.2g min elapsed'% ((time.time() - t0)/60.)

    Modelsdictz = cPickle.load(file(dictz, 'rb'))

    data.DICTS(filtersz, Modelsdictz)

    P = parspace.Pdict (data)  # Dictionary with all parameter space especifications.
                                # From PARAMETERSPACE_AGNfitter.py

    t1= time.time()

    MCMC_AGNfitter.main(data, P, mc)
    PLOTandWRITE_AGNfitter.main(data,  P,  out)

    print '_____________________________________________________'
    print 'For this fit %.2g min elapsed'% ((time.time() - t1)/60.)
    return

def RUN_AGNfitter_onesource( line, data_obj, modelsdict):
    """
    Main function for fitting a single source in line 'line'.
    """

    mc = MCMC_settings()
    out = OUTPUT_settings()

    data = DATA(data_obj,line)
    data.DICTS(filters, Modelsdict)

    P = parspace.Pdict (data)  # Dictionary with all parameter space especifications.
                                # From PARAMETERSPACE_AGNfitter.py

    print ''
    print 'Fitting sources from catalog: ', data.catalog
    print '- Sourceline: ', line
    print '- Sourcename: ', data.name

    t1= time.time()

    MCMC_AGNfitter.main(data, P, mc)
    PLOTandWRITE_AGNfitter.main(data,  P,  out)

    print '_____________________________________________________'
    print 'For this fit %.2g min elapsed'% ((time.time() - t1)/60.)
    return

def RUN_AGNfitter_multiprocessing(processors, data_obj, modelsdict):
    """
    Main function for fitting all sources in a large catalog.
    Splits the job of running the large number of sources
    into a chosen number of processors.
    """
    cat = CATALOG_settings()

    nsources = data_ALL.cat['nsources']

    pool = mp.Pool(processes = processors)
    catalog_fitting = pool.map(RUN_AGNfitter_onesource,range(nsources),data_obj, modelsdict)
    pool.close()
    pool.join()
    ##WRITE ALL RESULST IN ONE TABLE
    return

if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    parser.add_argument("-c","--ncpu", type=int, default=1, help="number of cpus to use for multiprocessing")
    parser.add_argument("-n", "--sourcenumber", type=int, default=-1, help="specify a single source number to run (this is the line number in hte catalogue not the source id/name)")
    parser.add_argument("-i","--independent", action="store_true", help="run independently per source, i.e. do not create a global model dictionary")

    args = parser.parse_args()

    cat = CATALOG_settings()
    filters= FILTERS_settings()
    data_ALL = DATA_all(cat)
    data_ALL.PROPS()

    if not os.path.isdir(os.path.abspath(cat['dict_path'])):
	os.system('mkdir -p '+os.path.abspath(cat['dict_path']))

    # run for once source only and construct dictionary only for this source
    if args.independent:
        RUN_AGNfitter_onesource_independent(args.sourcenumber, data_ALL)

    else:
        Modelsdict = MAKE_model_dictionary(cat)

    # a single source is specified
    if args.sourcenumber >= 0:
        RUN_AGNfitter_onesource(args.sourcenumber, data_ALL, Modelsdict)
    else:
        RUN_AGNfitter_multiprocessing(args.ncpu, data_ALL, Modelsdict)

    print '======= : ======='
    print 'Process finished.'
<commit_msg>tidy up<commit_after>
#!/usr/bin/env python2.7

"""%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

AGNfitter

==============================

Fitting SEDs of AGN in a MCMC Approach
G.Calistro Rivera, E.Lusso, J.Hennawi, D. Hogg

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

This is the main script.

For default use
change only the functions which state
***USER INPUT NEEDED***.

"""

#PYTHON IMPORTS
import sys,os
import numpy as np
import matplotlib.pyplot as plt
import pylab as pl
import time
import shelve
import multiprocessing as mp
import cPickle
import argparse

#AGNfitter IMPORTS
from functions import  MCMC_AGNfitter, PLOTandWRITE_AGNfitter
import functions.PARAMETERSPACE_AGNfitter as parspace
from functions.DATA_AGNfitter import DATA, DATA_all
from functions.PLOTandWRITE_AGNfitter import CHAIN, FLUXES_ARRAYS
from functions.DICTIONARIES_AGNfitter import MODELSDICT
import functions.CONSTRUCT_modelobjects as MODELFILES
from astropy import units as u

def CATALOG_settings():

    """==================================
    ***USER INPUT NEEDED***

    Set the right values to be able to read your catalog's format.
    FITS option is not available yet.
    =================================="""

    cat = dict()

    ##GENERAL
    cat['path'] ='/Users/USER/AGNfitter/'  #path to the AGNfitter code

    cat['filename'] = 'data/catalog_example.txt'
    cat['filetype'] = 'ASCII' ## catalog file type: 'ASCII' or 'FITS'.
                              ## FITS option not available yet.

    cat['name'] = 0#'ID'            ## If ASCII: Column index (int) of source IDs
                                        ## If FITS: not yet
    cat['redshift'] = 1#'z'              ## If ASCII:  Column index(int) of redshift
                                        ## If FITS: not yet

    ##FREQUENCIES/WAVELENGTHS
    ## if ASCII specify 'freq/wl_list', if FITS specify 'freq/wl_suffix'
    cat['freq/wl_list'] = np.arange(2,56,3).tolist()
                                        ## If ASCII: List of column indexes (int),
                                        ##           corresponding to freq/wl.
    #cat['freq/wl_suffix'] = '_wl'      ## If FITS: common ending to wavelength column names
    cat['freq/wl_format'] = 'wavelength' ## Gives your catalog *observed*
                                         ## 'frequency' or 'wavelength'?
    cat['freq/wl_unit'] = u.Angstrom       ## Astropy unit of freq or wavelength

    ##FLUXES
    ## if ASCII specify 'freq/wl_list', if FITS specify 'freq/wl_suffix'
    cat['flux_unit'] = u.Jy             ## Astropy unit of *flux* (astropy-units)
    cat['flux_list'] = np.arange(3,57,3).tolist()
                                        ## If ASCII: List of column indexes (int)
    #cat['flux_suffix'] = '_f'          ## If FITS: Common ending of all flux column names (str)

    cat['fluxerr_list'] = np.arange(4,58,3).tolist()
                                        ## If ASCII: List of column indexes (int)
    #cat['fluxerr_suffix'] = '_e'       ## If FITS: common ending to fluxerr column names (str)

    ##NON-DETECTIONS
    cat['ndflag_bool'] = False          ## Does you catalog has columns with flags 1(0) for
                                        ## detections (nondetections)?
    cat['ndflag_list'] = 'list'         ## If ASCII: List of column indexes (int)
                                        ## If FITS: List of column names (str)

    ##OTHERS (user input ***not*** needed)
    cat['outpath'] ='/car-data/wwilliams/bootes/multi_bootes/agnfitter/'   # allow for an output path outside of the AGNfitter code path
    cat['output_folder'] =  cat['outpath'] +'OUTPUT/'#if no special OUTPUT folder, leave default
    cat['dict_path'] = cat['outpath'] + 'models/MODELSDICT_default'

    return cat

def FILTERS_settings():

    """==================================
    Set the photometric bands included in your catalog,
    in order to integrate the models over their response curves.
    =================================="""

    filters = dict()

    filters['dict_zarray'] =np.array([0.283, 1.58])  # The grid of redshifts needed to fit your catalog
    filters['Bandset'] = 'BANDSET_default' # OPTIONS:
                                           # 'BANDSET_default' (for testing)
                                           # 'BANDSET_settings' (choosing relevant filters below, as given by your catalog)
                                           # if your filter is not included, go to DICTIONARIES_AGNfitter to add.

    filters['SPIRE500']= True
    filters['SPIRE350']= True
    filters['SPIRE250']= True
    filters['PACS160']=False
    filters['PACS100']=False

    filters['MIPS160']=False
    filters['MIPS70']=False
    filters['MIPS24']=True

    filters['IRAC4']=True
    filters['IRAC3']=True
    filters['IRAC2']=True
    filters['IRAC1']=True

    filters['WISE4']=False
    filters['WISE3']=False
    filters['WISE2']=False
    filters['WISE1']=False

    filters['Ks_2mass']=True
    filters['H_2mass']=True
    filters['J_2mass']=True

    filters['H_VISTA']=False
    filters['J_VISTA']=False
    filters['K_VISTA']=False
    filters['Y_VISTA']=True

    filters['u_SDSS']=False
    filters['g_SDSS']=False
    filters['r_SDSS']=False
    filters['i_SDSS']=False
    filters['z_SDSS']=False

    filters['g_SUBARU']=False
    filters['r_SUBARU']=True
    filters['i_SUBARU']=True
    filters['z_SUBARU']=True
    filters['B_SUBARU']=True
    filters['V_SUBARU']=False

    filters['u_CHFT']=True
    filters['g_CHFT']=False
    filters['r_CHFT']=False
    filters['i_CHFT']=False
    filters['z_CHFT']=False

    filters['GALEX_2500']=True
    filters['GALEX_1500']=False

    return filters

def MCMC_settings():

    """==================================
    Set your preferences for the MCMC sampling.
    =================================="""

    mc = dict()

    mc['Nwalkers'] = 100  ## number of walkers
    mc['Nburnsets']= 2   ## number of burn-in sets
    mc['Nburn'] = 4000 ## length of each burn-in sets
    mc['Nmcmc'] = 10000  ## length of each burn-in sets
    mc['iprint'] = 1000 ## show progress in terminal in steps of this many samples

    return mc

def OUTPUT_settings():

    """==================================
    Set your preferences for the production of OUTPUT files.
    =================================="""

    out = dict()

    out['plot_format'] = 'pdf'

    #CHAIN TRACES
    out['plot_tracesburn-in'] = False
    out['plot_tracesmcmc'] = True

    #BASIC OUTPUT
    out['Nsample'] = 1000 ## out['Nsample'] * out['Nthinning'] <= out['Nmcmc']
    out['Nthinning'] = 10 ## This describes thinning of the chain to sample
    out['writepar_meanwitherrors'] = True ##Write output values for all parameters in a file.
    out['plot_posteriortriangle'] = False ##Plot triangle with all parameters' PDFs?

    #INTEGRATED LUMINOSITIES
    out['calc_intlum'] = True
    out['realizations2int'] = 100 #This process is very time consuming.
                                #Around 100-1000 is recomendend for computational reasons.
                                #If you want to plot posterior triangles of
                                #the integrated luminosities, should be > 1000.
    out['plot_posteriortrianglewithluminosities'] = False  # requires out['calc_intlum']=True

    #INTEGRATION RANGES
    out['intlum_models'] = ['sb','bbb', 'bbbdered', 'gal', 'tor','sb']  #leave 'sb' always
                                                                        #as first element
    out['intlum_freqranges_unit'] = u.micron   #Astropy unit
    out['intlum_freqranges'] = np.array([[8.,1000.],[0.1,1.],[0.1,1.],[0.1,1.],[1.,30.],[1.,30.]])
    out['intlum_names'] = ['LIR(8-1000)','Lbb(0.1-1)', 'Lbbdered(0.1-1)', 'Lga(01-1)', 'Ltor(1-30)','Lsb(1-30)']

    #SED PLOTTING
    out['realizations2plot'] = 10

    out['plotSEDrealizations'] = True

    return out

def header():
    print '              '
    print '             XXXX'
    print '___________ XXX _________________________________________________'
    print '            XX      '
    print '            X     '
    print '            X                       AGNfitter                     '
    print '         __ X __                    ---------                '
    print '     /**\   |   /**\                                          '
    print '... (*** =  o  = ***) ...........................................'
    print '     \**/__ | __\**/                                     '
    print '            X              Fitting SEDs of AGN and Galaxies  '
    print '            X             in a MCMC Approach '
    print '           xx              (Calistro Rivera et al. 2016)    '
    print '          xx               '
    print '_______ xxx______________________________________________________'
    print '     xxxx'
    print ''
    return

def MAKE_model_dictionary(cat, clobbermodel=False):
    ## 0. CONSTRUCT DICTIONARY (not needed if default is used)
    t0= time.time()

    if clobbermodel and os.path.lexists(cat['dict_path']):
        print "removing model dictionary "+cat['dict_path']
        os.system('rm -rf '+ cat['dict_path'])

    if not os.path.lexists(cat['dict_path']):

        MODELFILES.construct(cat['path'])

        mydict = MODELSDICT(cat['dict_path'], cat['path'], filters)
        mydict.build()

        print '_____________________________________________________'
        print 'For this dictionary creation %.2g min elapsed'% ((time.time() - t0)/60.)

    Modelsdict = cPickle.load(file(cat['dict_path'], 'rb'))

    return Modelsdict

def RUN_AGNfitter_onesource_independent( line, data_obj, clobbermodel=False):
    """
    Main function for fitting a single source in line 'line' and create it's modelsdict independently.
    """

    mc = MCMC_settings()
    out = OUTPUT_settings()

    data = DATA(data_obj,line)
    data.DICTS(filters, Modelsdict)

    print ''
    print 'Fitting sources from catalog: ', data.catalog
    print '- Sourceline: ', line
    print '- Sourcename: ', data.name

    ## 0. CONSTRUCT DICTIONARY for this redshift
    t0= time.time()

    # needs a list/array of z
    filtersz = FILTERS_settings([data.z])

    # add a suffix for this source dictionary
    dictz = cat['dict_path'] + '_' + str(data.name)
    if clobbermodel and os.path.lexists(dictz):
        os.system('rm -rf '+dictz)
        print "removing source model dictionary "+dictz

    if not os.path.lexists(dictz):
        zdict = MODELSDICT(dictz, cat['path'], filtersz)
        zdict.build()
        print '_____________________________________________________'
        print 'For this dictionary creation %.2g min elapsed'% ((time.time() - t0)/60.)

    Modelsdictz = cPickle.load(file(dictz, 'rb'))

    data.DICTS(filtersz, Modelsdictz)

    P = parspace.Pdict (data)   # Dictionary with all parameter space especifications.
                                # From PARAMETERSPACE_AGNfitter.py

    t1= time.time()

    MCMC_AGNfitter.main(data, P, mc)
    PLOTandWRITE_AGNfitter.main(data,  P,  out)

    print '_____________________________________________________'
    print 'For this fit %.2g min elapsed'% ((time.time() - t1)/60.)
    return

def RUN_AGNfitter_onesource( line, data_obj, modelsdict):
    """
    Main function for fitting a single source in line 'line'.
    """

    mc = MCMC_settings()
    out = OUTPUT_settings()

    data = DATA(data_obj,line)
    data.DICTS(filters, Modelsdict)

    P = parspace.Pdict (data)  # Dictionary with all parameter space especifications.
                                # From PARAMETERSPACE_AGNfitter.py

    print ''
    print 'Fitting sources from catalog: ', data.catalog
    print '- Sourceline: ', line
    print '- Sourcename: ', data.name

    t1= time.time()

    MCMC_AGNfitter.main(data, P, mc)
    PLOTandWRITE_AGNfitter.main(data,  P,  out)

    print '_____________________________________________________'
    print 'For this fit %.2g min elapsed'% ((time.time() - t1)/60.)
    return

def RUN_AGNfitter_multiprocessing(processors, data_obj, modelsdict):
    """
    Main function for fitting all sources in a large catalog.
    Splits the job of running the large number of sources
    into a chosen number of processors.
    """
    cat = CATALOG_settings()

    nsources = data_ALL.cat['nsources']

    pool = mp.Pool(processes = processors)
    catalog_fitting = pool.map(RUN_AGNfitter_onesource,range(nsources),data_obj, modelsdict)
    pool.close()
    pool.join()
    ##WRITE ALL RESULST IN ONE TABLE
    return

if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    parser.add_argument("-c","--ncpu", type=int, default=1, help="number of cpus to use for multiprocessing")
    parser.add_argument("-n", "--sourcenumber", type=int, default=-1, help="specify a single source number to run (this is the line number in hte catalogue not the source id/name)")
    parser.add_argument("-i","--independent", action="store_true", help="run independently per source, i.e. do not create a global model dictionary")

    args = parser.parse_args()

    cat = CATALOG_settings()
    filters= FILTERS_settings()
    data_ALL = DATA_all(cat)
    data_ALL.PROPS()

    ## make sure the output paths exist
    if not os.path.isdir(os.path.abspath(cat['output_folder'])):
        os.system('mkdir -p '+os.path.abspath(cat['output_folder']))
    # abspath is needed because 'dict_path' is a file
    if not os.path.isdir(os.path.abspath(cat['dict_path'])):
        os.system('mkdir -p '+os.path.abspath(cat['dict_path']))

    # run for once source only and construct dictionary only for this source
    if args.independent:
        RUN_AGNfitter_onesource_independent(args.sourcenumber, data_ALL)

    else:
        # make/read the model dictionary
        Modelsdict = MAKE_model_dictionary(cat)

        # a single source is specified
        if args.sourcenumber >= 0:
            RUN_AGNfitter_onesource(args.sourcenumber, data_ALL, Modelsdict)
        else:
            RUN_AGNfitter_multiprocessing(args.ncpu, data_ALL, Modelsdict)

    print '======= : ======='
    print 'Process finished.'
<|endoftext|>


<commit_before>from util import elapsed
from time import time
from collections import defaultdict

import requests
import json
import re
from threading import Thread

def call_orcid_api(url):
    headers = {'Accept': 'application/orcid+json'}
    start = time()
    r = requests.get(url, headers=headers)
    print "got ORCID results in {elapsed}s for {url}".format(
        url=url,
        elapsed=elapsed(start)
    )
    orcid_resp_dict = r.json()
    return orcid_resp_dict

def get_orcid_api_raw_profile(id):
    url = "https://pub.orcid.org/v1.2/{id}/orcid-profile".format(id=id)
    orcid_resp_dict = call_orcid_api(url)
    return orcid_resp_dict["orcid-profile"]

def get_orcid_api_raw_affiliations(id):
    url = "https://pub.orcid.org/v1.2/{id}/affiliations".format(id=id)
    orcid_resp_dict = call_orcid_api(url)
    return orcid_resp_dict

def get_orcid_api_raw_funding(id):
    url = "https://pub.orcid.org/v1.2/{id}/funding".format(id=id)
    orcid_resp_dict = call_orcid_api(url)
    return orcid_resp_dict

# main constructor
def make_and_populate_orcid_profile(orcid_id):
    new_profile = OrcidProfile(orcid_id)
    new_profile.populate_from_orcid()
    return new_profile

# uses multithreaded approach from http://www.shanelynn.ie/using-python-threading-for-multiple-results-queue/
def build_and_return_orcid_profile(orcid_id, populated_profiles):
    print "building profile for", orcid_id
    new_profile = make_and_populate_orcid_profile(orcid_id)
    populated_profiles[orcid_id] = new_profile
    print "done building profile for", orcid_id
    return populated_profiles

def make_and_populate_all_orcid_profiles(orcid_ids):
    threads = []
    populated_profiles = {}

    # start a thread for each orcid_id. results stored in populated_profiles.
    for orcid_id in orcid_ids:
        process = Thread(target=build_and_return_orcid_profile, args=[orcid_id, populated_profiles])
        process.start()
        threads.append(process)

    # wait till all work is done
    for process in threads:
        process.join()

    # return the results
    orcid_profile_list = populated_profiles.values()
    return orcid_profile_list

def search_orcid(given_names, family_name):
    url = u"https://orcid.org/v1.2/search/orcid-bio/?q=given-names%3A{given_names}%20AND%20family-name%3A{family_name}&rows=100".format(
        given_names=given_names,
        family_name=family_name
    )
    orcid_resp_dict = call_orcid_api(url)
    ret = []
    try:
        orcid_results = orcid_resp_dict["orcid-search-results"]["orcid-search-result"]
    except TypeError:
        return ret

    orcid_ids = []
    for result in orcid_results:
        orcid_id = result["orcid-profile"]["orcid-identifier"]["path"]
        orcid_ids.append(orcid_id)

    orcid_profile_list = make_and_populate_all_orcid_profiles(orcid_ids)

    for orcid_profile in orcid_profile_list:
        # before dumping the object to a dictionary,
        # update this attribute based on the name we used to search for the orcid profile
        orcid_profile.has_name_variant_beyond_search_query = orcid_profile.has_only_this_name(given_names, family_name)
        orcid_profile_dict = orcid_profile.to_dict()
        ret.append(orcid_profile_dict)

    return sorted(ret, key=lambda k: k['sort_score'], reverse=True)

def get_most_recently_ended_activity(activities):
    if not activities:
        return None

    # if there are any without end dates, return them
    best = None
    for activity in activities:
        if not activity["end_year"]:
            best = activity
    if best:
        return best

    # else return the one that has most recently ended
    return sorted(activities, key=lambda k: k['end_year'], reverse=True)[0]

class OrcidProfile(object):
    def __init__(self, id):
        self.id = id
        # initialize this to False.  Will be overwritten by calling search code if it
        # determines the
        self.has_name_variant_beyond_search_query = False

    def populate_from_orcid(self):
        self.api_raw_profile = get_orcid_api_raw_profile(self.id)
        self.api_raw_affilations = get_orcid_api_raw_affiliations(self.id)
        self.api_raw_funding = get_orcid_api_raw_funding(self.id)

    @property
    def given_names(self):
        try:
            return self.api_raw_profile["orcid-bio"]["personal-details"]["given-names"]["value"]
        except (KeyError, TypeError):
            return None

    @property
    def family_name(self):
        try:
            return self.api_raw_profile["orcid-bio"]["personal-details"]["family-name"]["value"]
        except (KeyError, TypeError):
            return None

    @property
    def credit_name(self):
        try:
            return self.api_raw_profile["orcid-bio"]["personal-details"]["credit-name"]["value"]
        except (KeyError, TypeError):
            return None

    @property
    def other_names(self):
        try:
            return self.api_raw_profile["orcid-bio"]["personal-details"]["other-names"]["value"]
        except (KeyError, TypeError):
            return None

    def has_only_this_name(self, given_names, family_name):
        given_names = given_names.lower().strip()
        family_name = family_name.lower().strip()

        full_name = u"{} {}".format(given_names, family_name)

        if self.given_names.lower().strip() != given_names:
            return True
        if self.family_name.lower().strip() != family_name:
            return True
        if self.credit_name and self.credit_name.lower().strip() != full_name:
            return True
        if self.other_names and self.other_names.lower().strip() != full_name:
            return True
        return False

    @property
    def keywords(self):
        try:
            return self.api_raw_profile["orcid-bio"]["keywords"]["keyword"][0]["value"]
        except (KeyError, TypeError):
            return None

    @property
    def works(self):
        try:
            works = self.api_raw_profile["orcid-activities"]["orcid-works"]["orcid-work"]
        except TypeError:
            works = None

        if not works:
            works = []
        return works

    @property
    def latest_work(self):
        if not self.works:
            return None

        latest_year = 0
        for work in self.works:
            try:
                year = int(work["publication-date"]["year"]["value"])
            except TypeError:
                # no year found
                continue

            if year > latest_year:
                best = work
                latest_year = year

        if not best:
            best = self.works[0]

        return best["work-title"]["title"]["value"]

    @property
    def best_affiliation(self):
        return get_most_recently_ended_activity(self.affiliations)

    @property
    def affiliations(self):
        ret = []
        try:
            affiliation_list = self.api_raw_affilations["orcid-profile"]["orcid-activities"]["affiliations"]["affiliation"]
        except (TypeError, ):
            return ret

        for affl in affiliation_list:
            affl_name = affl["organization"]["name"]
            role_title = affl["role-title"]
            try:
                end_year = int(affl["end-date"]["year"]["value"])
            except (KeyError, TypeError,):
                end_year = None
            ret.append({
                "name": affl_name,
                "role_title": role_title,
                "end_year": end_year
                })
        return ret

    @property
    def best_funding(self):
        return get_most_recently_ended_activity(self.funding)

    @property
    def funding(self):
        ret = []
        try:
            funding_list = self.api_raw_funding["orcid-profile"]["orcid-activities"]["funding-list"]["funding"]
        except (TypeError, ):
            return ret

        for fund in funding_list:
            funder = fund["organization"]["name"]
            title = fund["funding-title"]["title"]["value"]
            try:
                end_year = int(fund["end-date"]["year"]["value"])
            except (KeyError, TypeError,):
                end_year = None
            ret.append({
                "funder": funder,
                "title": title,
                "end_year": end_year
                })
        return ret

    def __repr__(self):
        return u'<OrcidProfile ({id}) "{given_names} {family_name}" >'.format(
            id=self.id,
            given_names=self.given_names,
            family_name=self.family_name
        )

    def to_dict(self):
        search_clues_dict = {
            "credit_name": self.credit_name,
            "other_names": self.other_names,
            "keywords": self.keywords,
            "best_funding": self.best_funding,
            "best_affiliation": self.best_affiliation,
            "latest_work": self.latest_work,
            "has_name_variant_beyond_search_query": self.has_name_variant_beyond_search_query,
            "num_works": len(self.works)
        }

        search_clues = []
        for k, v in search_clues_dict.iteritems():
            search_clues.append({"key":k, "value":v})

        ret = {
            "id": self.id,
            "given_names": self.given_names,
            "family_name": self.family_name,
            "search_clues_list": search_clues,
            "sort_score": len([val for val in search_clues_dict.values() if val])
        }

        return ret

<commit_msg>don't need other orcid api calls after all!<commit_after>from util import elapsed
from time import time
from collections import defaultdict

import requests
import json
import re
from threading import Thread

def call_orcid_api(url):
    headers = {'Accept': 'application/orcid+json'}
    start = time()
    r = requests.get(url, headers=headers)
    print "got ORCID results in {elapsed}s for {url}".format(
        url=url,
        elapsed=elapsed(start)
    )
    orcid_resp_dict = r.json()
    return orcid_resp_dict

def get_orcid_api_raw_profile(id):
    url = "https://pub.orcid.org/v1.2/{id}/orcid-profile".format(id=id)
    orcid_resp_dict = call_orcid_api(url)
    return orcid_resp_dict["orcid-profile"]

# main constructor
def make_and_populate_orcid_profile(orcid_id):
    new_profile = OrcidProfile(orcid_id)
    new_profile.populate_from_orcid()
    return new_profile

# uses multithreaded approach from http://www.shanelynn.ie/using-python-threading-for-multiple-results-queue/
def build_and_return_orcid_profile(orcid_id, populated_profiles):
    print "building profile for", orcid_id
    new_profile = make_and_populate_orcid_profile(orcid_id)
    populated_profiles[orcid_id] = new_profile
    print "done building profile for", orcid_id
    return populated_profiles

def make_and_populate_all_orcid_profiles(orcid_ids):
    threads = []
    populated_profiles = {}

    # start a thread for each orcid_id. results stored in populated_profiles.
    for orcid_id in orcid_ids:
        process = Thread(target=build_and_return_orcid_profile, args=[orcid_id, populated_profiles])
        process.start()
        threads.append(process)

    # wait till all work is done
    for process in threads:
        process.join()

    # return the results
    orcid_profile_list = populated_profiles.values()
    return orcid_profile_list

def search_orcid(given_names, family_name):
    url = u"https://orcid.org/v1.2/search/orcid-bio/?q=given-names%3A{given_names}%20AND%20family-name%3A{family_name}&rows=100".format(
        given_names=given_names,
        family_name=family_name
    )
    orcid_resp_dict = call_orcid_api(url)
    ret = []
    try:
        orcid_results = orcid_resp_dict["orcid-search-results"]["orcid-search-result"]
    except TypeError:
        return ret

    orcid_ids = []
    for result in orcid_results:
        orcid_id = result["orcid-profile"]["orcid-identifier"]["path"]
        orcid_ids.append(orcid_id)

    orcid_profile_list = make_and_populate_all_orcid_profiles(orcid_ids)

    for orcid_profile in orcid_profile_list:
        # before dumping the object to a dictionary,
        # update this attribute based on the name we used to search for the orcid profile
        orcid_profile.has_name_variant_beyond_search_query = orcid_profile.has_only_this_name(given_names, family_name)
        orcid_profile_dict = orcid_profile.to_dict()
        ret.append(orcid_profile_dict)

    return sorted(ret, key=lambda k: k['sort_score'], reverse=True)

def get_most_recently_ended_activity(activities):
    if not activities:
        return None

    # if there are any without end dates, return them
    best = None
    for activity in activities:
        if not activity["end_year"]:
            best = activity
    if best:
        return best

    # else return the one that has most recently ended
    return sorted(activities, key=lambda k: k['end_year'], reverse=True)[0]

class OrcidProfile(object):
    def __init__(self, id):
        self.id = id
        # initialize this to False.  Will be overwritten by calling search code if it
        # determines the
        self.has_name_variant_beyond_search_query = False

    def populate_from_orcid(self):
        self.api_raw_profile = get_orcid_api_raw_profile(self.id)

    @property
    def given_names(self):
        try:
            return self.api_raw_profile["orcid-bio"]["personal-details"]["given-names"]["value"]
        except (KeyError, TypeError):
            return None

    @property
    def family_name(self):
        try:
            return self.api_raw_profile["orcid-bio"]["personal-details"]["family-name"]["value"]
        except (KeyError, TypeError):
            return None

    @property
    def credit_name(self):
        try:
            return self.api_raw_profile["orcid-bio"]["personal-details"]["credit-name"]["value"]
        except (KeyError, TypeError):
            return None

    @property
    def other_names(self):
        try:
            return self.api_raw_profile["orcid-bio"]["personal-details"]["other-names"]["value"]
        except (KeyError, TypeError):
            return None

    def has_only_this_name(self, given_names, family_name):
        given_names = given_names.lower().strip()
        family_name = family_name.lower().strip()

        full_name = u"{} {}".format(given_names, family_name)

        if self.given_names.lower().strip() != given_names:
            return True
        if self.family_name.lower().strip() != family_name:
            return True
        if self.credit_name and self.credit_name.lower().strip() != full_name:
            return True
        if self.other_names and self.other_names.lower().strip() != full_name:
            return True
        return False

    @property
    def biography(self):
        try:
            return self.api_raw_profile["orcid-bio"]["biography"]["value"]
        except (KeyError, TypeError):
            return None

    @property
    def keywords(self):
        try:
            return self.api_raw_profile["orcid-bio"]["keywords"]["keyword"][0]["value"]
        except (KeyError, TypeError):
            return None

    @property
    def works(self):
        try:
            works = self.api_raw_profile["orcid-activities"]["orcid-works"]["orcid-work"]
        except TypeError:
            works = None

        if not works:
            works = []
        return works

    @property
    def latest_work(self):
        if not self.works:
            return None

        latest_year = 0
        for work in self.works:
            try:
                year = int(work["publication-date"]["year"]["value"])
            except TypeError:
                # no year found
                continue

            if year > latest_year:
                best = work
                latest_year = year

        if not best:
            best = self.works[0]

        return best["work-title"]["title"]["value"]

    @property
    def best_affiliation(self):
        return get_most_recently_ended_activity(self.affiliations)

    @property
    def affiliations(self):
        ret = []
        try:
            affiliation_list = self.api_raw_profile["orcid-activities"]["affiliations"]["affiliation"]
        except (TypeError, ):
            return ret

        for affl in affiliation_list:
            affl_name = affl["organization"]["name"]
            role_title = affl["role-title"]
            try:
                end_year = int(affl["end-date"]["year"]["value"])
            except (KeyError, TypeError,):
                end_year = None
            ret.append({
                "name": affl_name,
                "role_title": role_title,
                "end_year": end_year
                })
        return ret

    @property
    def best_funding(self):
        return get_most_recently_ended_activity(self.funding)

    @property
    def funding(self):
        ret = []
        try:
            funding_list = self.api_raw_profile["orcid-activities"]["funding-list"]["funding"]
        except (TypeError, ):
            return ret

        for fund in funding_list:
            funder = fund["organization"]["name"]
            title = fund["funding-title"]["title"]["value"]
            try:
                end_year = int(fund["end-date"]["year"]["value"])
            except (KeyError, TypeError,):
                end_year = None
            ret.append({
                "funder": funder,
                "title": title,
                "end_year": end_year
                })
        return ret

    def __repr__(self):
        return u'<OrcidProfile ({id}) "{given_names} {family_name}" >'.format(
            id=self.id,
            given_names=self.given_names,
            family_name=self.family_name
        )

    def to_dict(self):
        search_clues_dict = {
            "credit_name": self.credit_name,
            "other_names": self.other_names,
            "keywords": self.keywords,
            "best_funding": self.best_funding,
            "best_affiliation": self.best_affiliation,
            "latest_work": self.latest_work,
            "has_name_variant_beyond_search_query": self.has_name_variant_beyond_search_query,
            "num_works": len(self.works)
        }

        search_clues = []
        for k, v in search_clues_dict.iteritems():
            search_clues.append({"key":k, "value":v})

        ret = {
            "id": self.id,
            "given_names": self.given_names,
            "family_name": self.family_name,
            "search_clues_list": search_clues,
            "sort_score": len([val for val in search_clues_dict.values() if val])
        }

        return ret

<|endoftext|>